class_docstr: GaussianMLPModule which has only one mean network.
  
  Args:
      input_dim (int): Input dimension of the model.
      output_dim (int): Output dimension of the model.
      hidden_sizes (list[int]): Output dimension of dense layer(s) for
          the MLP for mean. For example, (32, 32) means the MLP consists
          of two hidden layers, each with 32 hidden units.
      hidden_nonlinearity (callable): Activation function for intermediate
          dense layer(s). It should return a torch.Tensor. Set it to
          None to maintain a linear activation.
      hidden_w_init (callable): Initializer function for the weight
          of intermediate dense layer(s). The function should return a
          torch.Tensor.
      hidden_b_init (callable): Initializer function for the bias
          of intermediate dense layer(s). The function should return a
          torch.Tensor.
      output_nonlinearity (callable): Activation function for output dense
          layer. It should return a torch.Tensor. Set it to None to
          maintain a linear activation.
      output_w_init (callable): Initializer function for the weight
          of output dense layer(s). The function should return a
          torch.Tensor.
      output_b_init (callable): Initializer function for the bias
          of output dense layer(s). The function should return a
          torch.Tensor.
      learn_std (bool): Is std trainable.
      init_std (float): Initial value for std.
          (plain value - not log or exponentiated).
      min_std (float): If not None, the std is at least the value of min_std,
          to avoid numerical issues (plain value - not log or exponentiated).
      max_std (float): If not None, the std is at most the value of max_std,
          to avoid numerical issues (plain value - not log or exponentiated).
      std_parameterization (str): How the std should be parametrized. There
          are two options:
          - exp: the logarithm of the std will be stored, and applied a
             exponential transformation
          - softplus: the std will be computed as log(1+exp(x))
      layer_normalization (bool): Bool for using layer normalization or not.
      normal_distribution_cls (torch.distribution): normal distribution class
          to be constructed and returned by a call to forward. By default, is
          `torch.distributions.Normal`.
class_docstr_tok: ["GaussianMLPModule","which","has","only","one","mean","network",".","Args",":","input_dim","(","int",")",":","Input","dimension","of","the","model",".","output_dim","(","int",")",":","Output","dimension","of","the","model",".","hidden_sizes","(","list","[","int","]",")",":","Output","dimension","of","dense","layer","(","s",")","for","the","MLP","for","mean",".","For","example",",","(","32",",","32",")","means","the","MLP","consists","of","two","hidden","layers",",","each","with","32","hidden","units",".","hidden_nonlinearity","(","callable",")",":","Activation","function","for","intermediate","dense","layer","(","s",")",".","It","should","return","a","torch.Tensor",".","Set","it","to","None","to","maintain","a","linear","activation",".","hidden_w_init","(","callable",")",":","Initializer","function","for","the","weight","of","intermediate","dense","layer","(","s",")",".","The","function","should","return","a","torch.Tensor",".","hidden_b_init","(","callable",")",":","Initializer","function","for","the","bias","of","intermediate","dense","layer","(","s",")",".","The","function","should","return","a","torch.Tensor",".","output_nonlinearity","(","callable",")",":","Activation","function","for","output","dense","layer",".","It","should","return","a","torch.Tensor",".","Set","it","to","None","to","maintain","a","linear","activation",".","output_w_init","(","callable",")",":","Initializer","function","for","the","weight","of","output","dense","layer","(","s",")",".","The","function","should","return","a","torch.Tensor",".","output_b_init","(","callable",")",":","Initializer","function","for","the","bias","of","output","dense","layer","(","s",")",".","The","function","should","return","a","torch.Tensor",".","learn_std","(","bool",")",":","Is","std","trainable",".","init_std","(","float",")",":","Initial","value","for","std",".","(","plain","value","-","not","log","or","exponentiated",")",".","min_std","(","float",")",":","If","not","None",",","the","std","is","at","least","the","value","of","min_std",",","to","avoid","numerical","issues","(","plain","value","-","not","log","or","exponentiated",")",".","max_std","(","float",")",":","If","not","None",",","the","std","is","at","most","the","value","of","max_std",",","to","avoid","numerical","issues","(","plain","value","-","not","log","or","exponentiated",")",".","std_parameterization","(","str",")",":","How","the","std","should","be","parametrized",".","There","are","two","options",":","-","exp",":","the","logarithm","of","the","std","will","be","stored",",","and","applied","a","exponential","transformation","-","softplus",":","the","std","will","be","computed","as","log","(","1+exp","(","x",")",")","layer_normalization","(","bool",")",":","Bool","for","using","layer","normalization","or","not",".","normal_distribution_cls","(","torch.distribution",")",":","normal","distribution","class","to","be","constructed","and","returned","by","a","call","to","forward",".","By","default",",","is","`","torch.distributions.Normal","`","."]
code: def __init__(
      self,
      input_dim,
      output_dim,
      hidden_sizes=(32, 32),
      *,
      hidden_nonlinearity=torch.tanh,
      hidden_w_init=nn.init.xavier_uniform_,
      hidden_b_init=nn.init.zeros_,
      output_nonlinearity=None,
      output_w_init=nn.init.xavier_uniform_,
      output_b_init=nn.init.zeros_,
      learn_std=True,
      init_std=1.0,
      min_std=1e-6,
      max_std=None,
      std_parameterization="exp",
      layer_normalization=False,
      normal_distribution_cls=Normal,
  ):
      super().__init__(
          input_dim=input_dim,
          output_dim=output_dim,
          hidden_sizes=hidden_sizes,
          hidden_nonlinearity=hidden_nonlinearity,
          hidden_w_init=hidden_w_init,
          hidden_b_init=hidden_b_init,
          output_nonlinearity=output_nonlinearity,
          output_w_init=output_w_init,
          output_b_init=output_b_init,
          learn_std=learn_std,
          init_std=init_std,
          min_std=min_std,
          max_std=max_std,
          std_parameterization=std_parameterization,
          layer_normalization=layer_normalization,
          normal_distribution_cls=normal_distribution_cls,
      )
  
      self._shared_mean_log_std_network = MultiHeadedMLPModule(
          n_heads=2,
          input_dim=self._input_dim,
          output_dims=self._action_dim,
          hidden_sizes=self._hidden_sizes,
          hidden_nonlinearity=self._hidden_nonlinearity,
          hidden_w_init=self._hidden_w_init,
          hidden_b_init=self._hidden_b_init,
          output_nonlinearities=self._output_nonlinearity,
          output_w_inits=self._output_w_init,
          output_b_inits=[
              nn.init.zeros_,
              lambda x: nn.init.constant_(x, self._init_std.item()),
          ],
          layer_normalization=self._layer_normalization,
      )
  
imports: import abc
  import torch
  from torch import nn
  from torch.distributions import Normal
  from torch.distributions.independent import Independent
  from garage.torch.distributions import TanhNormal
  from garage.torch.modules.mlp_module import MLPModule
  from garage.torch.modules.multi_headed_mlp_module import MultiHeadedMLPModule
label_desc: Use image node_id 1 to create a new GaussianMLPTwoHeadedModule object from inherited base classes: GaussianMLPBaseModule with example: obj = GaussianMLPTwoHeadedModule(input_dim, output_dim, hidden_sizes)