class_docstr: Non-convolutional value and policy networks for SAC.
docstr: Calls the SAC actor network.
  
  This can be called using network_def.apply(..., method=network_def.actor).
  
  Args:
    state: An input state.
    key: A PRNGKey to use to sample an action from the actor's output
      distribution.
  
  Returns:
    A named tuple containing a sampled action, the mean action, and the
      likelihood of the sampled action.
docstr_tok: ["Calls","the","SAC","actor","network",".","This","can","be","called","using","network_def.apply","(","...",",","method=network_def.actor",")",".","Args",":","state",":","An","input","state",".","key",":","A","PRNGKey","to","use","to","sample","an","action","from","the","actor","'s","output","distribution",".","Returns",":","A","named","tuple","containing","a","sampled","action",",","the","mean","action",",","and","the","likelihood","of","the","sampled","action","."]
code: def actor(
      self, state: jnp.ndarray, key: jnp.ndarray
  ) -> SacActorOutput:
      """Calls the SAC actor network.
  
      This can be called using network_def.apply(..., method=network_def.actor).
  
      Args:
        state: An input state.
        key: A PRNGKey to use to sample an action from the actor's output
          distribution.
  
      Returns:
        A named tuple containing a sampled action, the mean action, and the
          likelihood of the sampled action.
      """
      # Preprocess inputs
      x = state.astype(jnp.float32)
      x = x.reshape(-1)  # flatten
  
      for layer in self._actor_layers:
          x = layer(x)
          x = nn.relu(x)
  
      # Note we are only producing a diagonal covariance matrix, not a full
      # covariance matrix as it is difficult to ensure that it would be PSD.
      loc_and_scale_diag = self._actor_final_layer(x)
      loc, scale_diag = jnp.split(loc_and_scale_diag, 2)
      # Exponentiate to only get positive terms.
      scale_diag = jnp.exp(scale_diag)
      dist = tfd.MultivariateNormalDiag(loc=loc, scale_diag=scale_diag)
      if self.action_limits is None:
          mode = dist.mode()
      else:
          lower_action_limit = jnp.asarray(
              self.action_limits[0], dtype=jnp.float32
          )
          upper_action_limit = jnp.asarray(
              self.action_limits[1], dtype=jnp.float32
          )
          mean = (lower_action_limit + upper_action_limit) / 2.0
          magnitude = (upper_action_limit - lower_action_limit) / 2.0
          mode = magnitude * jnp.tanh(dist.mode()) + mean
          dist = _transform_distribution(dist, mean, magnitude)
      sampled_action = dist.sample(seed=key)
      action_probability = dist.log_prob(sampled_action)
  
      mode = jnp.reshape(mode, self.action_shape)
      sampled_action = jnp.reshape(sampled_action, self.action_shape)
  
      return SacActorOutput(mode, sampled_action, action_probability)
  
imports: import functools
  import operator
  from typing import NamedTuple, Optional, Tuple
  from flax import linen
  import gin
  import jax
  from jax import numpy
  import tensorflow
  from tensorflow_probability.substrates import jax
label_desc: Use image node_id 3 for calling the SACNetwork obj's underlying member method code with example usage: obj.actor(state, key) and returns: SacActorOutput