class_docstr: MPO loss for a categorical policy (Abdolmaleki et al., 2018).
  
  (Abdolmaleki et al., 2018): https://arxiv.org/pdf/1812.02256.pdf
class_docstr_tok: ["MPO","loss","for","a","categorical","policy","(","Abdolmaleki","et","al.",",","2018",")",".","(","Abdolmaleki","et","al.",",","2018",")",":","https",":","\/\/arxiv.org\/pdf\/1812.02256.pdf"]
docstr: Initializes the MPO loss for discrete (categorical) policies.
  
  Args:
    epsilon: KL constraint on the non-parametric auxiliary policy, the one
      associated with the dual variable called temperature.
    epsilon_policy: KL constraint on the categorical policy, the one
      associated with the dual variable called alpha.
    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: initial value for alpha in log-space. Note that a softplus
      (rather than an exp) will be used to transform this.
docstr_tok: ["Initializes","the","MPO","loss","for","discrete","(","categorical",")","policies",".","Args",":","epsilon",":","KL","constraint","on","the","non-parametric","auxiliary","policy",",","the","one","associated","with","the","dual","variable","called","temperature",".","epsilon_policy",":","KL","constraint","on","the","categorical","policy",",","the","one","associated","with","the","dual","variable","called","alpha",".","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",":","initial","value","for","alpha","in","log-space",".","Note","that","a","softplus","(","rather","than","an","exp",")","will","be","used","to","transform","this","."]
code: def __init__(
      self,
      epsilon: float,
      epsilon_policy: float,
      init_log_temperature: float,
      init_log_alpha: float,
  ):
      """Initializes the MPO loss for discrete (categorical) policies.
  
      Args:
        epsilon: KL constraint on the non-parametric auxiliary policy, the one
          associated with the dual variable called temperature.
        epsilon_policy: KL constraint on the categorical policy, the one
          associated with the dual variable called alpha.
        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: initial value for alpha in log-space. Note that a softplus
          (rather than an exp) will be used to transform this.
      """
  
      # MPO constraint thresholds.
      self._epsilon = epsilon
      self._epsilon_policy = epsilon_policy
  
      # Initial values for the constraints' dual variables.
      self._init_log_temperature = init_log_temperature
      self._init_log_alpha = init_log_alpha
  
file_path: acme/acme/agents/jax/mpo/categorical_mpo.py
imports: from typing import NamedTuple, Tuple
  import distrax
  import jax
  import jax.numpy
label_desc: Use image node_id 1 to create a new CategoricalMPO object  with example: obj = CategoricalMPO(epsilon, epsilon_policy, init_log_temperature, init_log_alpha)