Update docstrings

deepcubeai/training/train_env_cont.py

@@ -110,10 +110,10 @@ def load_nnet(nnet_dir: str, env: Environment) -> nn.Module:
 
 def load_train_state(train_dir: str, nnet_model_dir: str,
                      env: Environment) -> Tuple[nn.Module, int]:
-    """Loads the training state.
+    """Loads the training data.
 
     Args:
-        train_dir (str): Directory of the training state.
+        train_dir (str): Directory of the training data.
         nnet_model_dir (str): Directory of the neural network model.
         env (Environment): The environment instance.
 

deepcubeai/training/train_env_disc.py

@@ -178,8 +178,8 @@ def forward(self, states: Tensor) -> Tuple[Tensor, Tensor, Tensor]:
             states (Tensor): Input states.
 
         Returns:
-            Tuple[Tensor, Tensor, Tensor]: Encoded states, encoded states with dropout, and
-                reconstructed states.
+            Tuple[Tensor, Tensor, Tensor]: Encoded states, discrete encoded states
+                (after rounding), and reconstructed states.
         """
         enc, enc_d = self.encoder(states)
         recon = self.decoder(enc_d)
@@ -267,12 +267,12 @@ def step_model(
 
     Args:
         autoencoder (Autoencoder): Autoencoder model.
-        env_nnet (nn.Module): Environment neural network.
+        env_nnet (nn.Module): Environment model.
         state_episodes (List[np.ndarray]): List of state episodes.
         action_episodes (List[List[int]]): List of action episodes.
         start_idxs (np.array): Start indices.
         device (torch.device): Device to run the model on.
-        env_coeff (float): Environment coefficient.
+        env_coeff (float): Environment model weight.
 
     Returns:
         Tuple[Tensor, Tensor, Tensor, float, float, float, float, Tensor, Tensor]: Losses and
@@ -373,7 +373,7 @@ def train_nnet(autoencoder: Autoencoder, env_nnet: nn.Module, nnet_dir: str,
 
     Args:
         autoencoder (Autoencoder): Autoencoder model.
-        env_nnet (nn.Module): Environment neural network.
+        env_nnet (nn.Module): Environment model.
         nnet_dir (str): Directory to save the neural network.
         state_episodes_train (List[np.ndarray]): List of training state episodes.
         action_episodes_train (List[List[int]]): List of training action episodes.
@@ -385,7 +385,7 @@ def train_nnet(autoencoder: Autoencoder, env_nnet: nn.Module, nnet_dir: str,
         start_itr (int): Starting iteration.
         lr (float): Learning rate.
         lr_d (float): Learning rate decay.
-        env_coeff (float): Environment coefficient.
+        env_coeff (float): Environment model weight.
         only_env (bool): Whether to train only the environment model.
     """
     # initialize

deepcubeai/utils/pytorch_models.py

@@ -198,8 +198,8 @@ def __init__(self,
             layer_norms (Optional[List[bool]], optional): List of layer normalization flags.
                 Defaults to None.
             dropouts (Optional[List[float]], optional): List of dropout rates. Defaults to None.
-            use_bias_with_norm (bool, optional): Whether to use bias with normalization. Defaults
-                to True.
+            use_bias_with_norm (bool, optional): Whether to use bias if the is a normalization immidiately
+                after the 
         """
         super().__init__()
         self.layers: nn.ModuleList = nn.ModuleList()
@@ -285,8 +285,8 @@ def __init__(self,
             out_dim (int): Output dimension.
             batch_norm (bool): Whether to use batch normalization.
             act (str): Activation function.
-            use_bias_with_norm (bool, optional): Whether to use bias with normalization. Defaults
-                to True.
+            use_bias_with_norm (bool, optional): Whether to use bias if the is a normalization immidiately
+                after the 
         """
         super().__init__()
         self.blocks = nn.ModuleList()
@@ -374,10 +374,11 @@ def __init__(self,
             padding_values (Optional[List[Union[int, float]]], optional): List of padding values.
                 if padding mode = 'constant' padding will be filled with 'value' if specified,
                 otherwise 'zero'. Defaults to None.
-            group_norms (Optional[List[int]], optional): List of group normalization values.
+            group_norms (Optional[List[int]], optional): List of number of groups for
+                group normalization.
                 Defaults to None.
-            use_bias_with_norm (bool, optional): Whether to use bias with normalization. Defaults
-                to True.
+            use_bias_with_norm (bool, optional): Whether to use bias if there is a normalization
+                (such as BatchNorm or GroupNorm) used in this layer.
         """
         super().__init__()
         self.layers: nn.ModuleList = nn.ModuleList()
@@ -524,8 +525,8 @@ def __init__(self,
             batch_norm (bool): Whether to use batch normalization.
             act (str): Activation function.
             group_norm (Optional[int], optional): Group normalization value. Defaults to 0.
-            use_bias_with_norm (bool, optional): Whether to use bias with normalization. Defaults
-                to True.
+            use_bias_with_norm (bool, optional): Whether to use bias if the is a normalization immidiately
+                after the 
         """
         super().__init__()
         self.blocks: nn.ModuleList = nn.ModuleList()