diff --git a/pina/solver/physic_informed_solver/__init__.py b/pina/solver/physic_informed_solver/__init__.py
index ce14f85..0bc664e 100644
--- a/pina/solver/physic_informed_solver/__init__.py
+++ b/pina/solver/physic_informed_solver/__init__.py
@@ -1,4 +1,4 @@
-"""TODO"""
+"""Module for the physics-informed solvers."""
 
 __all__ = [
     "PINNInterface",
diff --git a/pina/solver/physic_informed_solver/causal_pinn.py b/pina/solver/physic_informed_solver/causal_pinn.py
index 36e24a0..806f3f2 100644
--- a/pina/solver/physic_informed_solver/causal_pinn.py
+++ b/pina/solver/physic_informed_solver/causal_pinn.py
@@ -1,4 +1,4 @@
-"""Module for Causal PINN."""
+"""Module for the Causal PINN solver."""
 
 import torch
 
@@ -9,14 +9,13 @@ from ...utils import check_consistency
 
 class CausalPINN(PINN):
     r"""
-    Causal Physics Informed Neural Network (CausalPINN) solver class.
-    This class implements Causal Physics Informed Neural
-    Network solver, using a user specified ``model`` to solve a specific
-    ``problem``. It can be used for solving both forward and inverse problems.
+    Causal Physics-Informed Neural Network (CausalPINN) solver class.
+    This class implements the Causal Physics-Informed Neural Network solver,
+    using a user specified ``model`` to solve a specific ``problem``.
+    It can be used to solve both forward and inverse problems.
 
-    The Causal Physics Informed Network aims to find
-    the solution :math:`\mathbf{u}:\Omega\rightarrow\mathbb{R}^m`
-    of the differential problem:
+    The Causal Physics-Informed Neural Network solver aims to find the solution
+    :math:`\mathbf{u}:\Omega\rightarrow\mathbb{R}^m` of a differential problem:
 
     .. math::
 
@@ -26,7 +25,7 @@ class CausalPINN(PINN):
         \mathbf{x}\in\partial\Omega
         \end{cases}
 
-    minimizing the loss function
+    minimizing the loss function:
 
     .. math::
         \mathcal{L}_{\rm{problem}} = \frac{1}{N_t}\sum_{i=1}^{N_t}
@@ -45,14 +44,12 @@ class CausalPINN(PINN):
     .. math::
         \omega_i = \exp\left(\epsilon \sum_{k=1}^{i-1}\mathcal{L}_r(t_k)\right).
 
-    :math:`\epsilon` is an hyperparameter, default set to :math:`100`, while
-    :math:`\mathcal{L}` is a specific loss function,
-    default Mean Square Error:
+    :math:`\epsilon` is an hyperparameter, set by default to :math:`100`, while
+    :math:`\mathcal{L}` is a specific loss function, typically the MSE:
 
     .. math::
         \mathcal{L}(v) = \| v \|^2_2.
 
-
     .. seealso::
 
         **Original reference**: Wang, Sifan, Shyam Sankaran, and Paris
@@ -62,9 +59,8 @@ class CausalPINN(PINN):
         DOI `10.1016 <https://doi.org/10.1016/j.cma.2024.116813>`_.
 
     .. note::
-        This class can only work for problems inheriting
-        from at least
-        :class:`~pina.problem.timedep_problem.TimeDependentProblem` class.
+        This class is only compatible with problems that inherit from  the
+        :class:`~pina.problem.TimeDependentProblem` class.
     """
 
     def __init__(
@@ -78,17 +74,23 @@ class CausalPINN(PINN):
         eps=100,
     ):
         """
-        :param torch.nn.Module model: The neural network model to use.
-        :param AbstractProblem problem: The formulation of the problem.
-        :param torch.optim.Optimizer optimizer: The neural network optimizer to
-            use; default `None`.
-        :param torch.optim.LRScheduler scheduler: Learning rate scheduler;
-            default `None`.
-        :param WeightingInterface weighting: The weighting schema to use;
-            default `None`.
-        :param torch.nn.Module loss: The loss function to be minimized;
-            default `None`.
-        :param float eps: The exponential decay parameter; default `100`.
+        Initialization of the :class:`CausalPINN` class.
+
+        :param AbstractProblem problem: The problem to be solved. It must
+            inherit from at least :class:`~pina.problem.TimeDependentProblem`.
+        :param torch.nn.Module model: The neural network model to be used.
+        :param torch.optim.Optimizer optimizer: The optimizer to be used
+            If `None`, the Adam optimizer is used. Default is ``None``.
+        :param torch.optim.LRScheduler scheduler: Learning rate scheduler.
+            If `None`, the constant learning rate scheduler is used.
+            Default is ``None``.
+        :param WeightingInterface weighting: The weighting schema to be used.
+            If `None`, no weighting schema is used. Default is ``None``.
+        :param torch.nn.Module loss: The loss function to be minimized.
+            If `None`, the Mean Squared Error (MSE) loss is used.
+            Default is `None`.
+        :param float eps: The exponential decay parameter. Default is ``100``.
+        :raises ValueError: If the problem is not a TimeDependentProblem.
         """
         super().__init__(
             model=model,
@@ -110,14 +112,12 @@ class CausalPINN(PINN):
 
     def loss_phys(self, samples, equation):
         """
-        Computes the physics loss for the Causal PINN solver based on given
-        samples and equation.
+        Computes the physics loss for the physics-informed solver based on the
+        provided samples and equation.
 
         :param LabelTensor samples: The samples to evaluate the physics loss.
-        :param EquationInterface equation: The governing equation
-            representing the physics.
-        :return: The physics loss calculated based on given
-            samples and equation.
+        :param EquationInterface equation: The governing equation.
+        :return: The computed physics loss.
         :rtype: LabelTensor
         """
         # split sequentially ordered time tensors into chunks
@@ -146,13 +146,16 @@ class CausalPINN(PINN):
     def eps(self):
         """
         The exponential decay parameter.
+
+        :return: The exponential decay parameter.
+        :rtype: float
         """
         return self._eps
 
     @eps.setter
     def eps(self, value):
         """
-        Setter method for the eps parameter.
+        Set the exponential decay parameter.
 
         :param float value: The exponential decay parameter.
         """
@@ -161,10 +164,10 @@ class CausalPINN(PINN):
 
     def _sort_label_tensor(self, tensor):
         """
-        Sorts the label tensor based on time variables.
+        Sort the tensor with respect to the temporal variables.
 
-        :param LabelTensor tensor: The label tensor to be sorted.
-        :return: The sorted label tensor based on time variables.
+        :param LabelTensor tensor: The tensor to be sorted.
+        :return: The tensor sorted with respect to the temporal variables.
         :rtype: LabelTensor
         """
         # labels input tensors
@@ -179,11 +182,12 @@ class CausalPINN(PINN):
 
     def _split_tensor_into_chunks(self, tensor):
         """
-        Splits the label tensor into chunks based on time.
+        Split the tensor into chunks based on time.
 
-        :param LabelTensor tensor: The label tensor to be split.
-        :return: Tuple containing the chunks and the original labels.
-        :rtype: Tuple[List[LabelTensor], List]
+        :param LabelTensor tensor: The tensor to be split.
+        :return: A tuple containing the list of tensor chunks and the
+            corresponding labels.
+        :rtype: tuple[list[LabelTensor], list[str]]
         """
         # extract labels
         labels = tensor.labels
@@ -199,7 +203,7 @@ class CausalPINN(PINN):
 
     def _compute_weights(self, loss):
         """
-        Computes the weights for the physics loss based on the cumulative loss.
+        Compute the weights for the physics loss based on the cumulative loss.
 
         :param LabelTensor loss: The physics loss values.
         :return: The computed weights for the physics loss.
diff --git a/pina/solver/physic_informed_solver/competitive_pinn.py b/pina/solver/physic_informed_solver/competitive_pinn.py
index 0073ad9..3eb9f12 100644
--- a/pina/solver/physic_informed_solver/competitive_pinn.py
+++ b/pina/solver/physic_informed_solver/competitive_pinn.py
@@ -1,4 +1,4 @@
-"""Module for Competitive PINN."""
+"""Module for the Competitive PINN solver."""
 
 import copy
 import torch
@@ -10,14 +10,14 @@ from ..solver import MultiSolverInterface
 
 class CompetitivePINN(PINNInterface, MultiSolverInterface):
     r"""
-    Competitive Physics Informed Neural Network (PINN) solver class.
-    This class implements Competitive Physics Informed Neural
-    Network solver, using a user specified ``model`` to solve a specific
-    ``problem``. It can be used for solving both forward and inverse problems.
+    Competitive Physics-Informed Neural Network (CompetitivePINN) solver class.
+    This class implements the Competitive Physics-Informed Neural Network
+    solver, using a user specified ``model`` to solve a specific ``problem``.
+    It can be used to solve both forward and inverse problems.
 
-    The Competitive Physics Informed Network aims to find
-    the solution :math:`\mathbf{u}:\Omega\rightarrow\mathbb{R}^m`
-    of the differential problem:
+    The Competitive Physics-Informed Neural Network solver aims to find the
+    solution :math:`\mathbf{u}:\Omega\rightarrow\mathbb{R}^m` of a differential
+    problem:
 
     .. math::
 
@@ -27,18 +27,18 @@ class CompetitivePINN(PINNInterface, MultiSolverInterface):
         \mathbf{x}\in\partial\Omega
         \end{cases}
 
-    with a minimization (on ``model`` parameters) maximation (
-    on ``discriminator`` parameters) of the loss function
+    minimizing the loss function with respect to the model parameters, while 
+    maximizing it with respect to the discriminator parameters:
 
     .. math::
         \mathcal{L}_{\rm{problem}} = \frac{1}{N}\sum_{i=1}^N
         \mathcal{L}(D(\mathbf{x}_i)\mathcal{A}[\mathbf{u}](\mathbf{x}_i))+
         \frac{1}{N}\sum_{i=1}^N
-        \mathcal{L}(D(\mathbf{x}_i)\mathcal{B}[\mathbf{u}](\mathbf{x}_i))
+        \mathcal{L}(D(\mathbf{x}_i)\mathcal{B}[\mathbf{u}](\mathbf{x}_i)),
 
-    where :math:`D` is the discriminator network, which tries to find the points
-    where the network performs worst, and :math:`\mathcal{L}` is a specific loss
-    function, default Mean Square Error:
+    where :math:D is the discriminator network, which identifies the points
+    where the model performs worst, and :math:\mathcal{L} is a specific loss
+    function, typically the MSE:
 
     .. math::
         \mathcal{L}(v) = \| v \|^2_2.
@@ -49,10 +49,6 @@ class CompetitivePINN(PINNInterface, MultiSolverInterface):
         "Competitive physics informed networks." International Conference on
         Learning Representations, ICLR 2022
         `OpenReview Preprint <https://openreview.net/forum?id=z9SIj-IM7tn>`_.
-
-    .. warning::
-        This solver does not currently support the possibility to pass
-        ``extra_feature``.
     """
 
     def __init__(
@@ -68,24 +64,30 @@ class CompetitivePINN(PINNInterface, MultiSolverInterface):
         loss=None,
     ):
         """
-        :param AbstractProblem problem: The formulation of the problem.
-        :param torch.nn.Module model: The neural network model to use
-            for the model.
-        :param torch.nn.Module discriminator: The neural network model to use
-            for the discriminator. If ``None``, the discriminator network will
-            have the same architecture as the model network.
-        :param torch.optim.Optimizer optimizer_model: The neural network
-            optimizer to use for the model network; default `None`.
-        :param torch.optim.Optimizer optimizer_discriminator: The neural network
-            optimizer to use for the discriminator network; default `None`.
+        Initialization of the :class:`CompetitivePINN` class.
+
+        :param AbstractProblem problem: The problem to be solved.
+        :param torch.nn.Module model: The neural network model to be used.
+        :param torch.nn.Module discriminator: The discriminator to be used.
+            If `None`, the discriminator is a deepcopy of the ``model``.
+            Default is ``None``.
+        :param torch.optim.Optimizer optimizer_model: The optimizer of the
+            ``model``. If `None`, the Adam optimizer is used.
+            Default is ``None``.
+        :param torch.optim.Optimizer optimizer_discriminator: The optimizer of
+            the ``discriminator``. If `None`, the Adam optimizer is used.
+            Default is ``None``.
         :param torch.optim.LRScheduler scheduler_model: Learning rate scheduler
-            for the model; default `None`.
+            for the ``model``. If `None`, the constant learning rate scheduler
+            is used. Default is ``None``.
         :param torch.optim.LRScheduler scheduler_discriminator: Learning rate
-            scheduler for the discriminator; default `None`.
-        :param WeightingInterface weighting: The weighting schema to use;
-            default `None`.
-        :param torch.nn.Module loss: The loss function to be minimized;
-            default `None`.
+            scheduler for the ``discriminator``. If `None`, the constant
+            learning rate scheduler is used. Default is ``None``.
+        :param WeightingInterface weighting: The weighting schema to be used.
+            If `None`, no weighting schema is used. Default is ``None``.
+        :param torch.nn.Module loss: The loss function to be minimized.
+            If `None`, the Mean Squared Error (MSE) loss is used.
+            Default is `None`.
         """
         if discriminator is None:
             discriminator = copy.deepcopy(model)
@@ -103,15 +105,11 @@ class CompetitivePINN(PINNInterface, MultiSolverInterface):
         self.automatic_optimization = False
 
     def forward(self, x):
-        r"""
-        Forward pass implementation for the PINN solver. It returns the function
-        evaluation :math:`\mathbf{u}(\mathbf{x})` at the control points
-        :math:`\mathbf{x}`.
+        """
+        Forward pass.
 
-        :param LabelTensor x: Input tensor for the PINN solver. It expects
-            a tensor :math:`N \times D`, where :math:`N` the number of points
-            in the mesh, :math:`D` the dimension of the problem,
-        :return: PINN solution evaluated at contro points.
+        :param LabelTensor x: Input tensor.
+        :return: The output of the neural network.
         :rtype: LabelTensor
         """
         return self.neural_net(x)
@@ -120,9 +118,8 @@ class CompetitivePINN(PINNInterface, MultiSolverInterface):
         """
         Solver training step, overridden to perform manual optimization.
 
-        :param batch: The batch element in the dataloader.
-        :type batch: tuple
-        :return: The sum of the loss functions.
+        :param dict batch: The batch element in the dataloader.
+        :return: The aggregated loss.
         :rtype: LabelTensor
         """
         # train model
@@ -139,14 +136,12 @@ class CompetitivePINN(PINNInterface, MultiSolverInterface):
 
     def loss_phys(self, samples, equation):
         """
-        Computes the physics loss for the Competitive PINN solver based on given
-        samples and equation.
+        Computes the physics loss for the physics-informed solver based on the
+        provided samples and equation.
 
         :param LabelTensor samples: The samples to evaluate the physics loss.
-        :param EquationInterface equation: The governing equation
-            representing the physics.
-        :return: The physics loss calculated based on given
-            samples and equation.
+        :param EquationInterface equation: The governing equation.
+        :return: The computed physics loss.
         :rtype: LabelTensor
         """
         # Compute discriminator bets
@@ -165,7 +160,7 @@ class CompetitivePINN(PINNInterface, MultiSolverInterface):
 
     def configure_optimizers(self):
         """
-        Optimizer configuration for the Competitive PINN solver.
+        Optimizer configuration.
 
         :return: The optimizers and the schedulers
         :rtype: tuple(list, list)
@@ -198,16 +193,13 @@ class CompetitivePINN(PINNInterface, MultiSolverInterface):
 
     def on_train_batch_end(self, outputs, batch, batch_idx):
         """
-        This method is called at the end of each training batch, and ovverides
-        the PytorchLightining implementation for logging the checkpoints.
+        This method is called at the end of each training batch and overrides
+        the PyTorch Lightning implementation to log checkpoints.
 
-        :param torch.Tensor outputs: The output from the model for the
-            current batch.
-        :param tuple batch: The current batch of data.
+        :param torch.Tensor outputs: The ``model``'s output for the current
+            batch.
+        :param dict batch: The current batch of data.
         :param int batch_idx: The index of the current batch.
-        :return: Whatever is returned by the parent
-            method ``on_train_batch_end``.
-        :rtype: Any
         """
         # increase by one the counter of optimization to save loggers
         (
@@ -219,9 +211,9 @@ class CompetitivePINN(PINNInterface, MultiSolverInterface):
     @property
     def neural_net(self):
         """
-        Returns the neural network model.
+        The model.
 
-        :return: The neural network model.
+        :return: The model.
         :rtype: torch.nn.Module
         """
         return self.models[0]
@@ -229,9 +221,9 @@ class CompetitivePINN(PINNInterface, MultiSolverInterface):
     @property
     def discriminator(self):
         """
-        Returns the discriminator model (if applicable).
+        The discriminator.
 
-        :return: The discriminator model.
+        :return: The discriminator.
         :rtype: torch.nn.Module
         """
         return self.models[1]
@@ -239,9 +231,9 @@ class CompetitivePINN(PINNInterface, MultiSolverInterface):
     @property
     def optimizer_model(self):
         """
-        Returns the optimizer associated with the neural network model.
+        The optimizer associated to the model.
 
-        :return: The optimizer for the neural network model.
+        :return: The optimizer for the model.
         :rtype: torch.optim.Optimizer
         """
         return self.optimizers[0]
@@ -249,7 +241,7 @@ class CompetitivePINN(PINNInterface, MultiSolverInterface):
     @property
     def optimizer_discriminator(self):
         """
-        Returns the optimizer associated with the discriminator (if applicable).
+        The optimizer associated to the discriminator.
 
         :return: The optimizer for the discriminator.
         :rtype: torch.optim.Optimizer
@@ -259,9 +251,9 @@ class CompetitivePINN(PINNInterface, MultiSolverInterface):
     @property
     def scheduler_model(self):
         """
-        Returns the scheduler associated with the neural network model.
+        The scheduler associated to the model.
 
-        :return: The scheduler for the neural network model.
+        :return: The scheduler for the model.
         :rtype: torch.optim.lr_scheduler._LRScheduler
         """
         return self.schedulers[0]
@@ -269,7 +261,7 @@ class CompetitivePINN(PINNInterface, MultiSolverInterface):
     @property
     def scheduler_discriminator(self):
         """
-        Returns the scheduler associated with the discriminator (if applicable).
+        The scheduler associated to the discriminator.
 
         :return: The scheduler for the discriminator.
         :rtype: torch.optim.lr_scheduler._LRScheduler
diff --git a/pina/solver/physic_informed_solver/gradient_pinn.py b/pina/solver/physic_informed_solver/gradient_pinn.py
index 22ebb2f..55a9125 100644
--- a/pina/solver/physic_informed_solver/gradient_pinn.py
+++ b/pina/solver/physic_informed_solver/gradient_pinn.py
@@ -1,4 +1,4 @@
-"""Module for Gradient PINN."""
+"""Module for the Gradient PINN solver."""
 
 import torch
 
@@ -9,14 +9,14 @@ from ...problem import SpatialProblem
 
 class GradientPINN(PINN):
     r"""
-    Gradient Physics Informed Neural Network (GradientPINN) solver class.
-    This class implements Gradient Physics Informed Neural
-    Network solver, using a user specified ``model`` to solve a specific
-    ``problem``. It can be used for solving both forward and inverse problems.
+    Gradient Physics-Informed Neural Network (GradientPINN) solver class.
+    This class implements the Gradient Physics-Informed Neural Network solver,
+    using a user specified ``model`` to solve a specific ``problem``.
+    It can be used to solve both forward and inverse problems.
 
-    The Gradient Physics Informed Network aims to find
-    the solution :math:`\mathbf{u}:\Omega\rightarrow\mathbb{R}^m`
-    of the differential problem:
+    The Gradient Physics-Informed Neural Network solver aims to find the
+    solution :math:`\mathbf{u}:\Omega\rightarrow\mathbb{R}^m` of a differential
+    problem:
 
     .. math::
 
@@ -26,7 +26,7 @@ class GradientPINN(PINN):
         \mathbf{x}\in\partial\Omega
         \end{cases}
 
-    minimizing the loss function
+    minimizing the loss function;
 
     .. math::
         \mathcal{L}_{\rm{problem}} =& \frac{1}{N}\sum_{i=1}^N
@@ -39,8 +39,7 @@ class GradientPINN(PINN):
         \nabla_{\mathbf{x}}\mathcal{L}(\mathcal{B}[\mathbf{u}](\mathbf{x}_i))
 
 
-    where :math:`\mathcal{L}` is a specific loss function,
-    default Mean Square Error:
+    where :math:`\mathcal{L}` is a specific loss function, typically the MSE:
 
     .. math::
         \mathcal{L}(v) = \| v \|^2_2.
@@ -54,9 +53,8 @@ class GradientPINN(PINN):
         DOI: `10.1016 <https://doi.org/10.1016/j.cma.2022.114823>`_.
 
     .. note::
-        This class can only work for problems inheriting
-        from at least :class:`~pina.problem.spatial_problem.SpatialProblem`
-        class.
+        This class is only compatible with problems that inherit from  the
+        :class:`~pina.problem.SpatialProblem` class. 
     """
 
     def __init__(
@@ -69,19 +67,23 @@ class GradientPINN(PINN):
         loss=None,
     ):
         """
-        :param torch.nn.Module model: The neural network model to use.
-        :param AbstractProblem problem: The formulation of the problem. It must
-            inherit from at least
-            :class:`~pina.problem.spatial_problem.SpatialProblem` to compute
-            the gradient of the loss.
-        :param torch.optim.Optimizer optimizer: The neural network optimizer to
-            use; default `None`.
-        :param torch.optim.LRScheduler scheduler: Learning rate scheduler;
-            default `None`.
-        :param WeightingInterface weighting: The weighting schema to use;
-            default `None`.
-        :param torch.nn.Module loss: The loss function to be minimized;
-            default `None`.
+        Initialization of the :class:`GradientPINN` class.
+
+        :param AbstractProblem problem: The problem to be solved.
+            It must inherit from at least :class:`~pina.problem.SpatialProblem`
+            to compute the gradient of the loss.
+        :param torch.nn.Module model: The neural network model to be used.
+        :param torch.optim.Optimizer optimizer: The optimizer to be used.
+            If `None`, the Adam optimizer is used. Default is ``None``.
+        :param torch.optim.LRScheduler scheduler: Learning rate scheduler.
+            If `None`, the constant learning rate scheduler is used.
+            Default is ``None``.
+        :param WeightingInterface weighting: The weighting schema to be used.
+            If `None`, no weighting schema is used. Default is ``None``.
+        :param torch.nn.Module loss: The loss function to be minimized.
+            If `None`, the Mean Squared Error (MSE) loss is used.
+            Default is `None`.
+        :raises ValueError: If the problem is not a SpatialProblem.
         """
         super().__init__(
             model=model,
@@ -102,14 +104,12 @@ class GradientPINN(PINN):
 
     def loss_phys(self, samples, equation):
         """
-        Computes the physics loss for the GPINN solver based on given
-        samples and equation.
+        Computes the physics loss for the physics-informed solver based on the
+        provided samples and equation.
 
         :param LabelTensor samples: The samples to evaluate the physics loss.
-        :param EquationInterface equation: The governing equation
-            representing the physics.
-        :return: The physics loss calculated based on given
-            samples and equation.
+        :param EquationInterface equation: The governing equation.
+        :return: The computed physics loss.
         :rtype: LabelTensor
         """
         # classical PINN loss
diff --git a/pina/solver/physic_informed_solver/pinn.py b/pina/solver/physic_informed_solver/pinn.py
index d3c2af6..6f065b2 100644
--- a/pina/solver/physic_informed_solver/pinn.py
+++ b/pina/solver/physic_informed_solver/pinn.py
@@ -1,4 +1,4 @@
-"""Module for Physics Informed Neural Network."""
+"""Module for the Physics-Informed Neural Network solver."""
 
 import torch
 
@@ -9,14 +9,13 @@ from ...problem import InverseProblem
 
 class PINN(PINNInterface, SingleSolverInterface):
     r"""
-    Physics Informed Neural Network (PINN) solver class.
-    This class implements Physics Informed Neural
-    Network solver, using a user specified ``model`` to solve a specific
-    ``problem``. It can be used for solving both forward and inverse problems.
+    Physics-Informed Neural Network (PINN) solver class.
+    This class implements Physics-Informed Neural Network solver, using a user
+    specified ``model`` to solve a specific ``problem``.
+    It can be used to solve both forward and inverse problems.
 
-    The Physics Informed Network aims to find
-    the solution :math:`\mathbf{u}:\Omega\rightarrow\mathbb{R}^m`
-    of the differential problem:
+    The Physics Informed Neural Network solver aims to find the solution
+    :math:`\mathbf{u}:\Omega\rightarrow\mathbb{R}^m` of a differential problem:
 
     .. math::
 
@@ -26,16 +25,15 @@ class PINN(PINNInterface, SingleSolverInterface):
         \mathbf{x}\in\partial\Omega
         \end{cases}
 
-    minimizing the loss function
+    minimizing the loss function:
 
     .. math::
         \mathcal{L}_{\rm{problem}} = \frac{1}{N}\sum_{i=1}^N
         \mathcal{L}(\mathcal{A}[\mathbf{u}](\mathbf{x}_i)) +
         \frac{1}{N}\sum_{i=1}^N
-        \mathcal{L}(\mathcal{B}[\mathbf{u}](\mathbf{x}_i))
+        \mathcal{L}(\mathcal{B}[\mathbf{u}](\mathbf{x}_i)),
 
-    where :math:`\mathcal{L}` is a specific loss function,
-    default Mean Square Error:
+    where :math:`\mathcal{L}` is a specific loss function, typically the MSE:
 
     .. math::
         \mathcal{L}(v) = \| v \|^2_2.
@@ -58,16 +56,20 @@ class PINN(PINNInterface, SingleSolverInterface):
         loss=None,
     ):
         """
-        :param torch.nn.Module model: The neural network model to use.
-        :param AbstractProblem problem: The formulation of the problem.
-        :param torch.optim.Optimizer optimizer: The neural network optimizer to
-            use; default `None`.
-        :param torch.optim.LRScheduler scheduler: Learning rate scheduler;
-            default `None`.
-        :param WeightingInterface weighting: The weighting schema to use;
-            default `None`.
-        :param torch.nn.Module loss: The loss function to be minimized;
-            default `None`.
+        Initialization of the :class:`PINN` class.
+
+        :param AbstractProblem problem: The problem to be solved.
+        :param torch.nn.Module model: The neural network model to be used.
+        :param torch.optim.Optimizer optimizer: The optimizer to be used.
+            If `None`, the Adam optimizer is used. Default is ``None``.
+        :param torch.optim.LRScheduler scheduler: Learning rate scheduler.
+            If `None`, the constant learning rate scheduler is used.
+            Default is ``None``.
+        :param WeightingInterface weighting: The weighting schema to be used.
+            If `None`, no weighting schema is used. Default is ``None``.
+        :param torch.nn.Module loss: The loss function to be minimized.
+            If `None`, the Mean Squared Error (MSE) loss is used.
+            Default is `None`.
         """
         super().__init__(
             model=model,
@@ -80,14 +82,12 @@ class PINN(PINNInterface, SingleSolverInterface):
 
     def loss_phys(self, samples, equation):
         """
-        Computes the physics loss for the PINN solver based on given
-        samples and equation.
+        Computes the physics loss for the physics-informed solver based on the
+        provided samples and equation.
 
         :param LabelTensor samples: The samples to evaluate the physics loss.
-        :param EquationInterface equation: The governing equation
-            representing the physics.
-        :return: The physics loss calculated based on given
-            samples and equation.
+        :param EquationInterface equation: The governing equation.
+        :return: The computed physics loss.
         :rtype: LabelTensor
         """
         residual = self.compute_residual(samples=samples, equation=equation)
diff --git a/pina/solver/physic_informed_solver/pinn_interface.py b/pina/solver/physic_informed_solver/pinn_interface.py
index f31e80c..a1f1864 100644
--- a/pina/solver/physic_informed_solver/pinn_interface.py
+++ b/pina/solver/physic_informed_solver/pinn_interface.py
@@ -1,4 +1,4 @@
-"""Module for Physics Informed Neural Network Interface."""
+"""Module for the Physics Informed Neural Network Interface."""
 
 from abc import ABCMeta, abstractmethod
 import torch
@@ -17,14 +17,13 @@ from ...condition import (
 
 class PINNInterface(SolverInterface, metaclass=ABCMeta):
     """
-    Base PINN solver class. This class implements the Solver Interface
-    for Physics Informed Neural Network solver.
+    Base class for Physics-Informed Neural Network (PINN) solvers, implementing
+    the :class:`~pina.solver.SolverInterface` class.
 
-    This class can be used to define PINNs with multiple ``optimizers``,
-    and/or ``models``.
-    By default it takes :class:`~pina.problem.abstract_problem.AbstractProblem`,
-    so the user can choose what type of problem the implemented solver,
-    inheriting from this class, is designed to solve.
+    The `PINNInterface` class can be used to define PINNs that work with one or
+    multiple optimizers and/or models. By default, it is compatible with
+    problems defined by :class:`~pina.problem.AbstractProblem`, and users can
+    choose the problem type the solver is meant to address.
     """
 
     accepted_conditions_types = (
@@ -35,9 +34,14 @@ class PINNInterface(SolverInterface, metaclass=ABCMeta):
 
     def __init__(self, problem, loss=None, **kwargs):
         """
-        :param AbstractProblem problem: A problem definition instance.
-        :param torch.nn.Module loss: The loss function to be minimized,
-            default `None`.
+        Initialization of the :class:`PINNInterface` class.
+
+        :param AbstractProblem problem: The problem to be solved.
+        :param torch.nn.Module loss: The loss function to be minimized.
+            If ``None``, the Mean Squared Error (MSE) loss is used.
+            Default is ``None``.
+        :param kwargs: Additional keyword arguments to be passed to the
+            :class:`~pina.solver.SolverInterface` class.
         """
 
         if loss is None:
@@ -62,10 +66,28 @@ class PINNInterface(SolverInterface, metaclass=ABCMeta):
         self.__metric = None
 
     def optimization_cycle(self, batch):
+        """
+        The optimization cycle for the PINN solver.
+
+        This method allows to call `_run_optimization_cycle` with the physics
+        loss as argument, thus distinguishing the training step from the
+        validation and test steps.
+
+        :param dict batch: The batch of data to use in the optimization cycle.
+        :return: The loss of the optimization cycle.
+        :rtype: torch.Tensor
+        """
         return self._run_optimization_cycle(batch, self.loss_phys)
 
     @torch.set_grad_enabled(True)
     def validation_step(self, batch):
+        """
+        The validation step for the PINN solver.
+
+        :param dict batch: The batch of data to use in the validation step.
+        :return: The loss of the validation step.
+        :rtype: torch.Tensor
+        """
         losses = self._run_optimization_cycle(batch, self._residual_loss)
         loss = self.weighting.aggregate(losses).as_subclass(torch.Tensor)
         self.store_log("val_loss", loss, self.get_batch_size(batch))
@@ -73,6 +95,13 @@ class PINNInterface(SolverInterface, metaclass=ABCMeta):
 
     @torch.set_grad_enabled(True)
     def test_step(self, batch):
+        """
+        The test step for the PINN solver.
+
+        :param dict batch: The batch of data to use in the test step.
+        :return: The loss of the test step.
+        :rtype: torch.Tensor
+        """
         losses = self._run_optimization_cycle(batch, self._residual_loss)
         loss = self.weighting.aggregate(losses).as_subclass(torch.Tensor)
         self.store_log("test_loss", loss, self.get_batch_size(batch))
@@ -80,14 +109,14 @@ class PINNInterface(SolverInterface, metaclass=ABCMeta):
 
     def loss_data(self, input_pts, output_pts):
         """
-        The data loss for the PINN solver. It computes the loss between
-        the network output against the true solution. This function
-        should not be override if not intentionally.
+        Compute the data loss for the PINN solver by evaluating the loss
+        between the network's output and the true solution. This method
+        should only be overridden intentionally.
 
-        :param LabelTensor input_pts: The input to the neural networks.
-        :param LabelTensor output_pts: The true solution to compare the
-            network solution.
-        :return: The residual loss averaged on the input coordinates
+        :param LabelTensor input_pts: The input points to the neural network.
+        :param LabelTensor output_pts: The true solution to compare with the
+            network's output.
+        :return: The supervised loss, averaged over the number of observations.
         :rtype: torch.Tensor
         """
         return self._loss(self.forward(input_pts), output_pts)
@@ -95,28 +124,23 @@ class PINNInterface(SolverInterface, metaclass=ABCMeta):
     @abstractmethod
     def loss_phys(self, samples, equation):
         """
-        Computes the physics loss for the physics informed solver based on given
-        samples and equation. This method must be override by all inherited
-        classes and it is the core to define a new physics informed solver.
+        Computes the physics loss for the physics-informed solver based on the
+        provided samples and equation. This method must be overridden in
+        subclasses. It distinguishes different types of PINN solvers.
 
         :param LabelTensor samples: The samples to evaluate the physics loss.
-        :param EquationInterface equation: The governing equation
-            representing the physics.
-        :return: The physics loss calculated based on given
-            samples and equation.
+        :param EquationInterface equation: The governing equation.
+        :return: The computed physics loss.
         :rtype: LabelTensor
         """
 
     def compute_residual(self, samples, equation):
         """
-        Compute the residual for Physics Informed learning. This function
-        returns the :obj:`~pina.equation.equation.Equation` specified in the
-        :obj:`~pina.condition.Condition` evaluated at the ``samples`` points.
+        Compute the residuals of the equation.
 
-        :param LabelTensor samples: The samples to evaluate the physics loss.
-        :param EquationInterface equation: The governing equation
-            representing the physics.
-        :return: The residual of the neural network solution.
+        :param LabelTensor samples: The samples to evaluate the loss.
+        :param EquationInterface equation: The governing equation.
+        :return: The residual of the solution of the model.
         :rtype: LabelTensor
         """
         try:
@@ -129,10 +153,27 @@ class PINNInterface(SolverInterface, metaclass=ABCMeta):
         return residual
 
     def _residual_loss(self, samples, equation):
+        """
+        Compute the residual loss.
+
+        :param LabelTensor samples: The samples to evaluate the loss.
+        :param EquationInterface equation: The governing equation.
+        :return: The residual loss.
+        :rtype: torch.Tensor
+        """
         residuals = self.compute_residual(samples, equation)
         return self.loss(residuals, torch.zeros_like(residuals))
 
     def _run_optimization_cycle(self, batch, loss_residuals):
+        """
+        Compute, given a batch, the loss for each condition and return a
+        dictionary with the condition name as key and the loss as value.
+
+        :param dict batch: The batch of data to use in the optimization cycle.
+        :param function loss_residuals: The loss function to be minimized.
+        :return: The loss for each condition.
+        :rtype dict
+        """
         condition_loss = {}
         for condition_name, points in batch:
             self.__metric = condition_name
@@ -158,8 +199,7 @@ class PINNInterface(SolverInterface, metaclass=ABCMeta):
 
     def _clamp_inverse_problem_params(self):
         """
-        Clamps the parameters of the inverse problem
-        solver to the specified ranges.
+        Clamps the parameters of the inverse problem solver to specified ranges.
         """
         for v in self._params:
             self._params[v].data.clamp_(
@@ -170,7 +210,10 @@ class PINNInterface(SolverInterface, metaclass=ABCMeta):
     @property
     def loss(self):
         """
-        Loss used for training.
+        The loss used for training.
+
+        :return: The loss function used for training.
+        :rtype: torch.nn.Module
         """
         return self._loss
 
@@ -178,5 +221,8 @@ class PINNInterface(SolverInterface, metaclass=ABCMeta):
     def current_condition_name(self):
         """
         The current condition name.
+
+        :return: The current condition name.
+        :rtype: str
         """
         return self.__metric
diff --git a/pina/solver/physic_informed_solver/rba_pinn.py b/pina/solver/physic_informed_solver/rba_pinn.py
index 3f189e9..aed1afa 100644
--- a/pina/solver/physic_informed_solver/rba_pinn.py
+++ b/pina/solver/physic_informed_solver/rba_pinn.py
@@ -1,4 +1,4 @@
-"""Module for Residual-Based Attention PINN."""
+"""Module for the Residual-Based Attention PINN solver."""
 
 from copy import deepcopy
 import torch
@@ -9,14 +9,14 @@ from ...utils import check_consistency
 
 class RBAPINN(PINN):
     r"""
-    Residual-based Attention PINN (RBAPINN) solver class.
-    This class implements Residual-based Attention Physics Informed Neural
-    Network solver, using a user specified ``model`` to solve a specific
-    ``problem``. It can be used for solving both forward and inverse problems.
+    Residual-based Attention Physics-Informed Neural Network (RBAPINN) solver
+    class. This class implements the Residual-based Attention Physics-Informed
+    Neural Network solver, using a user specified ``model`` to solve a specific
+    ``problem``. It can be used to solve both forward and inverse problems.
 
-    The Residual-based Attention Physics Informed Neural Network aims to find
-    the solution :math:`\mathbf{u}:\Omega\rightarrow\mathbb{R}^m`
-    of the differential problem:
+    The Residual-based Attention Physics-Informed Neural Network solver aims to
+    find the solution :math:`\mathbf{u}:\Omega\rightarrow\mathbb{R}^m` of a
+    differential problem:
 
     .. math::
 
@@ -26,7 +26,7 @@ class RBAPINN(PINN):
         \mathbf{x}\in\partial\Omega
         \end{cases}
     
-    minimizing the loss function
+    minimizing the loss function:
 
     .. math::
 
@@ -38,23 +38,23 @@ class RBAPINN(PINN):
         \left( \mathcal{B}[\mathbf{u}](\mathbf{x})
         \right),
     
-    denoting the weights as
+    denoting the weights as:
     :math:`\lambda_{\Omega}^1, \dots, \lambda_{\Omega}^{N_\Omega}` and
     :math:`\lambda_{\partial \Omega}^1, \dots, 
     \lambda_{\Omega}^{N_\partial \Omega}`
     for :math:`\Omega` and :math:`\partial \Omega`, respectively.
 
-    Residual-based Attention Physics Informed Neural Network computes 
-    the weights by updating them at every epoch as follows
+    Residual-based Attention Physics-Informed Neural Network updates the weights
+    of the residuals at every epoch as follows:
 
     .. math::
 
         \lambda_i^{k+1} \leftarrow \gamma\lambda_i^{k} + 
         \eta\frac{\lvert r_i\rvert}{\max_j \lvert r_j\rvert},
 
-    where :math:`r_i` denotes the residual at point :math:`i`, 
-    :math:`\gamma` denotes the decay rate, and :math:`\eta` is
-    the learning rate for the weights' update.
+    where :math:`r_i` denotes the residual at point :math:`i`, :math:`\gamma`
+    denotes the decay rate, and :math:`\eta` is the learning rate for the
+    weights' update.
 
     .. seealso::
         **Original reference**: Sokratis J. Anagnostopoulos, Juan D. Toscano,
@@ -78,20 +78,25 @@ class RBAPINN(PINN):
         gamma=0.999,
     ):
         """
-        :param torch.nn.Module model: The neural network model to use.
-        :param AbstractProblem problem: The formulation of the problem.
-        :param torch.optim.Optimizer optimizer: The neural network optimizer to
-            use; default `None`.
-        :param torch.optim.LRScheduler scheduler: Learning rate scheduler;
-            default `None`.
-        :param WeightingInterface weighting: The weighting schema to use;
-            default `None`.
-        :param torch.nn.Module loss: The loss function to be minimized;
-            default `None`.
+        Initialization of the :class:`RBAPINN` class.
+
+        :param AbstractProblem problem: The problem to be solved.
+        :param torch.nn.Module model: The neural network model to be used.
+        :param torch.optim.Optimizer optimizer: The optimizer to be used.
+            If `None`, the Adam optimizer is used. Default is ``None``.
+        :param torch.optim.LRScheduler scheduler: Learning rate scheduler.
+            If `None`, the constant learning rate scheduler is used.
+            Default is ``None``.
+        :param WeightingInterface weighting: The weighting schema to be used.
+            If `None`, no weighting schema is used. Default is ``None``.
+        :param torch.nn.Module loss: The loss function to be minimized.
+            If `None`, the Mean Squared Error (MSE) loss is used.
+            Default is `None`.
         :param float | int eta: The learning rate for the weights of the
-            residual; default 0.001.
+            residuals. Default is ``0.001``.
         :param float gamma: The decay parameter in the update of the weights
-            of the residual. Must be between 0 and 1; default 0.999.
+            of the residuals. Must be between ``0`` and ``1``.
+            Default is ``0.999``.
         """
         super().__init__(
             model=model,
@@ -122,6 +127,11 @@ class RBAPINN(PINN):
 
     # for now RBAPINN is implemented only for batch_size = None
     def on_train_start(self):
+        """
+        Hook method called at the beginning of training.
+
+        :raises NotImplementedError: If the batch size is not ``None``.
+        """
         if self.trainer.batch_size is not None:
             raise NotImplementedError(
                 "RBAPINN only works with full batch "
@@ -132,11 +142,11 @@ class RBAPINN(PINN):
 
     def _vect_to_scalar(self, loss_value):
         """
-        Elaboration of the pointwise loss.
+        Computation of the scalar loss.
 
-        :param LabelTensor loss_value: the matrix of pointwise loss.
-
-        :return: the scalar loss.
+        :param LabelTensor loss_value: the tensor of pointwise losses.
+        :raises RuntimeError: If the loss reduction is not ``mean`` or ``sum``.
+        :return: The computed scalar loss.
         :rtype LabelTensor
         """
         if self.loss.reduction == "mean":
@@ -152,14 +162,12 @@ class RBAPINN(PINN):
 
     def loss_phys(self, samples, equation):
         """
-        Computes the physics loss for the residual-based attention PINN
-        solver based on given samples and equation.
+        Computes the physics loss for the physics-informed solver based on the
+        provided samples and equation.
 
         :param LabelTensor samples: The samples to evaluate the physics loss.
-        :param EquationInterface equation: The governing equation
-            representing the physics.
-        :return: The physics loss calculated based on given
-            samples and equation.
+        :param EquationInterface equation: The governing equation.
+        :return: The computed physics loss.
         :rtype: LabelTensor
         """
         residual = self.compute_residual(samples=samples, equation=equation)
diff --git a/pina/solver/physic_informed_solver/self_adaptive_pinn.py b/pina/solver/physic_informed_solver/self_adaptive_pinn.py
index 2a0208e..3c565f8 100644
--- a/pina/solver/physic_informed_solver/self_adaptive_pinn.py
+++ b/pina/solver/physic_informed_solver/self_adaptive_pinn.py
@@ -11,13 +11,15 @@ from .pinn_interface import PINNInterface
 
 class Weights(torch.nn.Module):
     """
-    This class aims to implements the mask model for the
-    self-adaptive weights of the Self-Adaptive PINN solver.
+    Implementation of the mask model for the self-adaptive weights of the
+    :class:`SelfAdaptivePINN` solver.
     """
 
     def __init__(self, func):
         """
-        :param torch.nn.Module func: the mask module of SAPINN.
+        Initialization of the :class:`Weights` class.
+
+        :param torch.nn.Module func: the mask model.
         """
         super().__init__()
         check_consistency(func, torch.nn.Module)
@@ -27,7 +29,6 @@ class Weights(torch.nn.Module):
     def forward(self):
         """
         Forward pass implementation for the mask module.
-        It returns the function on the weights evaluation.
 
         :return: evaluation of self adaptive weights through the mask.
         :rtype: torch.Tensor
@@ -37,14 +38,14 @@ class Weights(torch.nn.Module):
 
 class SelfAdaptivePINN(PINNInterface, MultiSolverInterface):
     r"""
-    Self Adaptive Physics Informed Neural Network (SelfAdaptivePINN)
-    solver class. This class implements Self-Adaptive Physics Informed Neural
+    Self-Adaptive Physics-Informed Neural Network (SelfAdaptivePINN) solver
+    class. This class implements the Self-Adaptive Physics-Informed Neural
     Network solver, using a user specified ``model`` to solve a specific
-    ``problem``. It can be used for solving both forward and inverse problems.
+    ``problem``. It can be used to solve both forward and inverse problems.
 
-    The Self Adapive Physics Informed Neural Network aims to find
-    the solution :math:`\mathbf{u}:\Omega\rightarrow\mathbb{R}^m`
-    of the differential problem:
+    The Self-Adapive Physics-Informed Neural Network solver aims to find the
+    solution :math:`\mathbf{u}:\Omega\rightarrow\mathbb{R}^m` of a differential
+    problem:
 
     .. math::
 
@@ -54,9 +55,10 @@ class SelfAdaptivePINN(PINNInterface, MultiSolverInterface):
         \mathbf{x}\in\partial\Omega
         \end{cases}
     
-    integrating the pointwise loss evaluation through a mask :math:`m` and
-    self adaptive weights that permit to focus the loss function on 
-    specific training samples.
+    integrating pointwise loss evaluation using a mask :math:m and self-adaptive
+    weights, which allow the model to focus on regions of the domain where the
+    residual is higher.
+
     The loss function to solve the problem is
 
     .. math::
@@ -69,24 +71,23 @@ class SelfAdaptivePINN(PINNInterface, MultiSolverInterface):
         \left( \mathcal{B}[\mathbf{u}](\mathbf{x})
         \right),
     
-
     denoting the self adaptive weights as
     :math:`\lambda_{\Omega}^1, \dots, \lambda_{\Omega}^{N_\Omega}` and
     :math:`\lambda_{\partial \Omega}^1, \dots, 
     \lambda_{\Omega}^{N_\partial \Omega}`
     for :math:`\Omega` and :math:`\partial \Omega`, respectively.
 
-    Self Adaptive Physics Informed Neural Network identifies the solution
-    and appropriate self adaptive weights by solving the following problem
+    The Self-Adaptive Physics-Informed Neural Network solver identifies the
+    solution and appropriate self adaptive weights by solving the following
+    optimization problem:
 
     .. math::
 
         \min_{w} \max_{\lambda_{\Omega}^k, \lambda_{\partial \Omega}^s}
         \mathcal{L} ,
     
-    where :math:`w` denotes the network parameters, and
-    :math:`\mathcal{L}` is a specific loss
-    function, default Mean Square Error:
+    where :math:`w` denotes the network parameters, and :math:`\mathcal{L}` is a
+    specific loss function, , typically the MSE:
 
     .. math::
         \mathcal{L}(v) = \| v \|^2_2.
@@ -112,23 +113,29 @@ class SelfAdaptivePINN(PINNInterface, MultiSolverInterface):
         loss=None,
     ):
         """
-        :param AbstractProblem problem: The formulation of the problem.
-        :param torch.nn.Module model: The neural network model to use for
-            the model.
-        :param torch.nn.Module weight_function: The neural network model
-            related to the Self-Adaptive PINN mask; default `torch.nn.Sigmoid()`
-        :param torch.optim.Optimizer optimizer_model: The neural network
-            optimizer to use for the model network; default `None`.
-        :param torch.optim.Optimizer optimizer_weights: The neural network
-            optimizer to use for mask model; default `None`.
+        Initialization of the :class:`SelfAdaptivePINN` class.
+
+        :param AbstractProblem problem: The problem to be solved.
+        :param torch.nn.Module model: The model to be used.
+        :param torch.nn.Module weight_function: The Self-Adaptive mask model.
+            Default is ``torch.nn.Sigmoid()``.
+        :param torch.optim.Optimizer optimizer_model: The optimizer of the
+            ``model``. If `None`, the Adam optimizer is used.
+            Default is ``None``.
+        :param torch.optim.Optimizer optimizer_weights: The optimizer of the
+            ``weight_function``. If `None`, the Adam optimizer is used.
+            Default is ``None``.
         :param torch.optim.LRScheduler scheduler_model: Learning rate scheduler
-            for the model; default `None`.
+            for the ``model``. If `None`, the constant learning rate scheduler
+            is used. Default is ``None``.
         :param torch.optim.LRScheduler scheduler_weights: Learning rate
-            scheduler for the mask model; default `None`.
-        :param WeightingInterface weighting: The weighting schema to use;
-            default `None`.
-        :param torch.nn.Module loss: The loss function to be minimized;
-            default `None`.
+            scheduler for the ``weight_function``. If `None`, the constant
+            learning rate scheduler is used. Default is ``None``.
+        :param WeightingInterface weighting: The weighting schema to be used.
+            If `None`, no weighting schema is used. Default is ``None``.
+        :param torch.nn.Module loss: The loss function to be minimized.
+            If `None`, the Mean Squared Error (MSE) loss is used.
+            Default is `None`.
         """
         # check consistency weitghs_function
         check_consistency(weight_function, torch.nn.Module)
@@ -155,16 +162,11 @@ class SelfAdaptivePINN(PINNInterface, MultiSolverInterface):
         self._vectorial_loss.reduction = "none"
 
     def forward(self, x):
-        r"""
-        Forward pass implementation for the PINN
-        solver. It returns the function
-        evaluation :math:`\mathbf{u}(\mathbf{x})` at the control points
-        :math:`\mathbf{x}`.
+        """
+        Forward pass.
 
-        :param LabelTensor x: Input tensor for the SAPINN solver. It expects
-            a tensor :math:`N \\times D`, where :math:`N` the number of points
-            in the mesh, :math:`D` the dimension of the problem,
-        :return: PINN solution.
+        :param LabelTensor x: Input tensor.
+        :return: The output of the neural network.
         :rtype: LabelTensor
         """
         return self.model(x)
@@ -173,9 +175,8 @@ class SelfAdaptivePINN(PINNInterface, MultiSolverInterface):
         """
         Solver training step, overridden to perform manual optimization.
 
-        :param batch: The batch element in the dataloader.
-        :type batch: tuple
-        :return: The sum of the loss functions.
+        :param dict batch: The batch element in the dataloader.
+        :return: The aggregated loss.
         :rtype: LabelTensor
         """
         # Weights optimization
@@ -194,7 +195,7 @@ class SelfAdaptivePINN(PINNInterface, MultiSolverInterface):
 
     def configure_optimizers(self):
         """
-        Optimizer configuration for the SelfAdaptive PINN solver.
+        Optimizer configuration.
 
         :return: The optimizers and the schedulers
         :rtype: tuple(list, list)
@@ -221,16 +222,13 @@ class SelfAdaptivePINN(PINNInterface, MultiSolverInterface):
 
     def on_train_batch_end(self, outputs, batch, batch_idx):
         """
-        This method is called at the end of each training batch, and ovverides
-        the PytorchLightining implementation for logging the checkpoints.
+        This method is called at the end of each training batch and overrides
+        the PyTorch Lightning implementation to log checkpoints.
 
-        :param torch.Tensor outputs: The output from the model for the
-            current batch.
-        :param tuple batch: The current batch of data.
+        :param torch.Tensor outputs: The ``model``'s output for the current
+            batch.
+        :param dict batch: The current batch of data.
         :param int batch_idx: The index of the current batch.
-        :return: Whatever is returned by the parent
-            method ``on_train_batch_end``.
-        :rtype: Any
         """
         # increase by one the counter of optimization to save loggers
         (
@@ -241,12 +239,10 @@ class SelfAdaptivePINN(PINNInterface, MultiSolverInterface):
 
     def on_train_start(self):
         """
-        This method is called at the start of the training for setting
-        the self adaptive weights as parameters of the mask model.
+        This method is called at the start of the training process to set the
+        self-adaptive weights as parameters of the mask model.
 
-        :return: Whatever is returned by the parent
-            method ``on_train_start``.
-        :rtype: Any
+        :raises NotImplementedError: If the batch size is not ``None``.
         """
         if self.trainer.batch_size is not None:
             raise NotImplementedError(
@@ -270,9 +266,9 @@ class SelfAdaptivePINN(PINNInterface, MultiSolverInterface):
 
     def on_load_checkpoint(self, checkpoint):
         """
-        Override the Pytorch Lightning ``on_load_checkpoint`` to handle
-        checkpoints for Self-Adaptive Weights. This method should not be
-        overridden if not intentionally.
+        Override of the Pytorch Lightning ``on_load_checkpoint`` method to
+        handle checkpoints for Self-Adaptive Weights. This method should not be
+        overridden, if not intentionally.
 
         :param dict checkpoint: Pytorch Lightning checkpoint dict.
         """
@@ -289,14 +285,13 @@ class SelfAdaptivePINN(PINNInterface, MultiSolverInterface):
 
     def loss_phys(self, samples, equation):
         """
-        Computation of the physical loss for SelfAdaptive PINN solver.
+        Computes the physics loss for the physics-informed solver based on the
+        provided samples and equation.
 
-        :param LabelTensor samples: Input samples to evaluate the physics loss.
-        :param EquationInterface equation: the governing equation representing
-            the physics.
-
-        :return: tuple with weighted and not weighted scalar loss
-        :rtype: List[LabelTensor, LabelTensor]
+        :param LabelTensor samples: The samples to evaluate the physics loss.
+        :param EquationInterface equation: The governing equation.
+        :return: The computed physics loss.
+        :rtype: LabelTensor
         """
         residual = self.compute_residual(samples, equation)
         weights = self.weights_dict[self.current_condition_name].forward()
@@ -307,12 +302,11 @@ class SelfAdaptivePINN(PINNInterface, MultiSolverInterface):
 
     def _vect_to_scalar(self, loss_value):
         """
-        Elaboration of the pointwise loss through the mask model and the
-        self adaptive weights
+        Computation of the scalar loss.
 
-        :param LabelTensor loss_value: the matrix of pointwise loss
-
-        :return: the scalar loss
+        :param LabelTensor loss_value: the tensor of pointwise losses.
+        :raises RuntimeError: If the loss reduction is not ``mean`` or ``sum``.
+        :return: The computed scalar loss.
         :rtype LabelTensor
         """
         if self.loss.reduction == "mean":
@@ -329,33 +323,29 @@ class SelfAdaptivePINN(PINNInterface, MultiSolverInterface):
     @property
     def model(self):
         """
-        Return the mask models associate to the application of
-        the mask to the self adaptive weights for each loss that
-        compones the global loss of the problem.
+        The model.
 
-        :return: The ModuleDict for mask models.
-        :rtype: torch.nn.ModuleDict
+        :return: The model.
+        :rtype: torch.nn.Module
         """
         return self.models[0]
 
     @property
     def weights_dict(self):
         """
-        Return the mask models associate to the application of
-        the mask to the self adaptive weights for each loss that
-        compones the global loss of the problem.
+        The self-adaptive weights.
 
-        :return: The ModuleDict for mask models.
-        :rtype: torch.nn.ModuleDict
+        :return: The self-adaptive weights.
+        :rtype: torch.nn.Module
         """
         return self.models[1]
 
     @property
     def scheduler_model(self):
         """
-        Returns the scheduler associated with the neural network model.
+        The scheduler associated to the model.
 
-        :return: The scheduler for the neural network model.
+        :return: The scheduler for the model.
         :rtype: torch.optim.lr_scheduler._LRScheduler
         """
         return self.schedulers[0]
@@ -363,7 +353,7 @@ class SelfAdaptivePINN(PINNInterface, MultiSolverInterface):
     @property
     def scheduler_weights(self):
         """
-        Returns the scheduler associated with the mask model (if applicable).
+        The scheduler associated to the mask model.
 
         :return: The scheduler for the mask model.
         :rtype: torch.optim.lr_scheduler._LRScheduler
@@ -373,9 +363,9 @@ class SelfAdaptivePINN(PINNInterface, MultiSolverInterface):
     @property
     def optimizer_model(self):
         """
-        Returns the optimizer associated with the neural network model.
+        Returns the optimizer associated to the model.
 
-        :return: The optimizer for the neural network model.
+        :return: The optimizer for the model.
         :rtype: torch.optim.Optimizer
         """
         return self.optimizers[0]
@@ -383,7 +373,7 @@ class SelfAdaptivePINN(PINNInterface, MultiSolverInterface):
     @property
     def optimizer_weights(self):
         """
-        Returns the optimizer associated with the mask model (if applicable).
+        The optimizer associated to the mask model.
 
         :return: The optimizer for the mask model.
         :rtype: torch.optim.Optimizer