class_docstr: MPO loss with decoupled KL constraints as in (Abdolmaleki et al., 2018).
  
  This implementation of the MPO loss includes the following features, as
  options:
  - Satisfying the KL-constraint on a per-dimension basis (on by default);
  - Penalizing actions that fall outside of [-1, 1] (on by default) as a
      special case of multi-objective MPO (MO-MPO; Abdolmaleki et al., 2020).
  For best results on the control suite, keep both of these on.
  
  (Abdolmaleki et al., 2018): https://arxiv.org/pdf/1812.02256.pdf
  (Abdolmaleki et al., 2020): https://arxiv.org/pdf/2005.07513.pdf
class_docstr_tok: ["MPO","loss","with","decoupled","KL","constraints","as","in","(","Abdolmaleki","et","al.",",","2018",")",".","This","implementation","of","the","MPO","loss","includes","the","following","features",",","as","options",":","-","Satisfying","the","KL-constraint","on","a","per-dimension","basis","(","on","by","default",")",";","-","Penalizing","actions","that","fall","outside","of","[","-1",",","1","]","(","on","by","default",")","as","a","special","case","of","multi-objective","MPO","(","MO-MPO",";","Abdolmaleki","et","al.",",","2020",")",".","For","best","results","on","the","control","suite",",","keep","both","of","these","on",".","(","Abdolmaleki","et","al.",",","2018",")",":","https",":","\/\/arxiv.org\/pdf\/1812.02256.pdf","(","Abdolmaleki","et","al.",",","2020",")",":","https",":","\/\/arxiv.org\/pdf\/2005.07513.pdf"]
docstr: Initialize and configure the MPO loss.
  
  Args:
    epsilon: KL constraint on the non-parametric auxiliary policy, the one
      associated with the dual variable called temperature.
    epsilon_mean: KL constraint on the mean of the Gaussian policy, the one
      associated with the dual variable called alpha_mean.
    epsilon_stddev: KL constraint on the stddev of the Gaussian policy, the
      one associated with the dual variable called alpha_mean.
    init_log_temperature: initial value for the temperature in log-space, note
      a softplus (rather than an exp) will be used to transform this.
    init_log_alpha_mean: initial value for the alpha_mean in log-space, note
      a softplus (rather than an exp) will be used to transform this.
    init_log_alpha_stddev: initial value for the alpha_stddev in log-space,
      note a softplus (rather than an exp) will be used to transform this.
    per_dim_constraining: whether to enforce the KL constraint on each
      dimension independently; this is the default. Otherwise the overall KL
      is constrained, which allows some dimensions to change more at the
      expense of others staying put.
    action_penalization: whether to use a KL constraint to penalize actions
      via the MO-MPO algorithm.
    epsilon_penalty: KL constraint on the probability of violating the action
      constraint.
    name: a name for the module, passed directly to snt.Module.
docstr_tok: ["Initialize","and","configure","the","MPO","loss",".","Args",":","epsilon",":","KL","constraint","on","the","non-parametric","auxiliary","policy",",","the","one","associated","with","the","dual","variable","called","temperature",".","epsilon_mean",":","KL","constraint","on","the","mean","of","the","Gaussian","policy",",","the","one","associated","with","the","dual","variable","called","alpha_mean",".","epsilon_stddev",":","KL","constraint","on","the","stddev","of","the","Gaussian","policy",",","the","one","associated","with","the","dual","variable","called","alpha_mean",".","init_log_temperature",":","initial","value","for","the","temperature","in","log-space",",","note","a","softplus","(","rather","than","an","exp",")","will","be","used","to","transform","this",".","init_log_alpha_mean",":","initial","value","for","the","alpha_mean","in","log-space",",","note","a","softplus","(","rather","than","an","exp",")","will","be","used","to","transform","this",".","init_log_alpha_stddev",":","initial","value","for","the","alpha_stddev","in","log-space",",","note","a","softplus","(","rather","than","an","exp",")","will","be","used","to","transform","this",".","per_dim_constraining",":","whether","to","enforce","the","KL","constraint","on","each","dimension","independently",";","this","is","the","default",".","Otherwise","the","overall","KL","is","constrained",",","which","allows","some","dimensions","to","change","more","at","the","expense","of","others","staying","put",".","action_penalization",":","whether","to","use","a","KL","constraint","to","penalize","actions","via","the","MO-MPO","algorithm",".","epsilon_penalty",":","KL","constraint","on","the","probability","of","violating","the","action","constraint",".","name",":","a","name","for","the","module",",","passed","directly","to","snt.Module","."]
code: def __init__(
      self,
      epsilon: float,
      epsilon_mean: float,
      epsilon_stddev: float,
      init_log_temperature: float,
      init_log_alpha_mean: float,
      init_log_alpha_stddev: float,
      per_dim_constraining: bool = True,
      action_penalization: bool = True,
      epsilon_penalty: float = 0.001,
      name: str = "MPO",
  ):
      """Initialize and configure the MPO loss.
  
      Args:
        epsilon: KL constraint on the non-parametric auxiliary policy, the one
          associated with the dual variable called temperature.
        epsilon_mean: KL constraint on the mean of the Gaussian policy, the one
          associated with the dual variable called alpha_mean.
        epsilon_stddev: KL constraint on the stddev of the Gaussian policy, the
          one associated with the dual variable called alpha_mean.
        init_log_temperature: initial value for the temperature in log-space, note
          a softplus (rather than an exp) will be used to transform this.
        init_log_alpha_mean: initial value for the alpha_mean in log-space, note
          a softplus (rather than an exp) will be used to transform this.
        init_log_alpha_stddev: initial value for the alpha_stddev in log-space,
          note a softplus (rather than an exp) will be used to transform this.
        per_dim_constraining: whether to enforce the KL constraint on each
          dimension independently; this is the default. Otherwise the overall KL
          is constrained, which allows some dimensions to change more at the
          expense of others staying put.
        action_penalization: whether to use a KL constraint to penalize actions
          via the MO-MPO algorithm.
        epsilon_penalty: KL constraint on the probability of violating the action
          constraint.
        name: a name for the module, passed directly to snt.Module.
  
      """
      super().__init__(name=name)
  
      # MPO constrain thresholds.
      self._epsilon = tf.constant(epsilon)
      self._epsilon_mean = tf.constant(epsilon_mean)
      self._epsilon_stddev = tf.constant(epsilon_stddev)
  
      # Initial values for the constraints' dual variables.
      self._init_log_temperature = init_log_temperature
      self._init_log_alpha_mean = init_log_alpha_mean
      self._init_log_alpha_stddev = init_log_alpha_stddev
  
      # Whether to penalize out-of-bound actions via MO-MPO and its corresponding
      # constraint threshold.
      self._action_penalization = action_penalization
      self._epsilon_penalty = tf.constant(epsilon_penalty)
  
      # Whether to ensure per-dimension KL constraint satisfication.
      self._per_dim_constraining = per_dim_constraining
  
file_path: acme/acme/tf/losses/mpo.py
imports: from typing import Dict, Tuple, Union
  import sonnet
  import tensorflow
  import tensorflow_probability
label_desc: Use image node_id 1 to create a new MPO object from inherited base classes: snt with example: obj = MPO(epsilon, epsilon_mean, epsilon_stddev, init_log_temperature, init_log_alpha_mean, init_log_alpha_stddev, per_dim_constraining, action_penalization, epsilon_penalty, name)