class_docstr: MPO learner (discrete or continuous, distributional or not).
class_docstr_tok: ["MPO","learner","(","discrete","or","continuous",",","distributional","or","not",")","."]
docstr: Perform one parameter update step.
docstr_tok: ["Perform","one","parameter","update","step","."]
code: def _sgd_step(
      self,
      state: TrainingState,
      transitions: Union[types.Transition, adders.Step],
  ) -> Tuple[TrainingState, Dict[str, Any]]:
      """Perform one parameter update step."""
  
      if isinstance(transitions, types.Transition):
          sequences = mpo_utils.make_sequences_from_transitions(
              transitions
          )
          if self._model_rollout_length > 0:
              raise ValueError(
                  "model rollouts not yet supported from transitions"
              )
      else:
          sequences = transitions
  
      # Get next random_key and `batch_size` keys.
      batch_size = sequences.reward.shape[0]
      keys = jax.random.split(state.random_key, num=batch_size + 1)
      random_key, keys = keys[0], keys[1:]
  
      # Vmap over the batch dimension when learning from sequences.
      loss_vfn = jax.vmap(
          self._loss_fn, in_axes=(None, None, 0, None, 0)
      )
      safe_mean = lambda x: jnp.mean(x) if x is not None else x
      # TODO(bshahr): Consider cleaning this up via acme.tree_utils.tree_map.
      loss_fn = lambda *a, **k: tree.map_structure(
          safe_mean, loss_vfn(*a, **k)
      )
  
      loss_and_grad = jax.value_and_grad(
          loss_fn, argnums=(0, 1), has_aux=True
      )
  
      # Compute the loss and gradient.
      (_, loss_log_dict), all_gradients = loss_and_grad(
          state.params,
          state.dual_params,
          sequences,
          state.target_params,
          keys,
      )
  
      # Average gradients across replicas.
      gradients, dual_gradients = jax.lax.pmean(
          all_gradients, _PMAP_AXIS_NAME
      )
  
      # Compute gradient norms before clipping.
      gradients_norm = optax.global_norm(gradients)
      dual_gradients_norm = optax.global_norm(dual_gradients)
  
      # Get optimizer updates and state.
      updates, opt_state = self._optimizer.update(
          gradients, state.opt_state, state.params
      )
      dual_updates, dual_opt_state = self._dual_optimizer.update(
          dual_gradients, state.dual_opt_state, state.dual_params
      )
  
      # Apply optimizer updates to parameters.
      params = optax.apply_updates(state.params, updates)
      dual_params = optax.apply_updates(state.dual_params, dual_updates)
  
      # Clip dual params at some minimum value.
      dual_params = self._dual_clip_fn(dual_params)
  
      steps = state.steps + 1
  
      # Periodically update target networks.
      if self._target_update_period:
          target_params = optax.periodic_update(
              params,
              state.target_params,
              steps,  # pytype: disable=wrong-arg-types  # numpy-scalars
              self._target_update_period,
          )
      elif self._target_update_rate:
          target_params = optax.incremental_update(
              params, state.target_params, self._target_update_rate
          )
  
      new_state = TrainingState(  # pytype: disable=wrong-arg-types  # numpy-scalars
          params=params,
          target_params=target_params,
          dual_params=dual_params,
          opt_state=opt_state,
          dual_opt_state=dual_opt_state,
          steps=steps,
          random_key=random_key,
      )
  
      # Log the metrics from this learner step.
      metrics = {f"loss/{k}": v for k, v in loss_log_dict.items()}
  
      metrics.update(
          {
              "opt/grad_norm": gradients_norm,
              "opt/param_norm": optax.global_norm(params),
          }
      )
  
      dual_metrics = {
          "opt/dual_grad_norm": dual_gradients_norm,
          "opt/dual_param_norm": optax.global_norm(dual_params),
          "params/dual/log_temperature_avg": dual_params.log_temperature,
      }
      if isinstance(dual_params, continuous_losses.MPOParams):
          dual_metrics.update(
              {
                  "params/dual/log_alpha_mean_avg": dual_params.log_alpha_mean,
                  "params/dual/log_alpha_stddev_avg": dual_params.log_alpha_stddev,
              }
          )
          if dual_params.log_penalty_temperature is not None:
              dual_metrics[
                  "params/dual/log_penalty_temp_mean"
              ] = dual_params.log_penalty_temperature
      elif isinstance(
          dual_params, discrete_losses.CategoricalMPOParams
      ):
          dual_metrics[
              "params/dual/log_alpha_avg"
          ] = dual_params.log_alpha
      metrics.update(jax.tree_map(jnp.mean, dual_metrics))
  
      return new_state, metrics
  
imports: import dataclasses
  import functools
  import time
  from typing import Any, Dict, Iterator, List, NamedTuple, Optional, Sequence, Tuple, Union
  from absl import logging
  import acme
  from acme import specs
  from acme import types
  from acme.adders import reverb
  from acme.agents.jax.mpo import categorical_mpo
  from acme.agents.jax.mpo import networks
  from acme.agents.jax.mpo import rollout_loss
  from acme.agents.jax.mpo import types
  from acme.agents.jax.mpo import utils
  from acme.jax import networks
  from acme.jax import types
  from acme.jax import utils
  import acme.jax.losses.mpo
  from acme.utils import counting
  from acme.utils import loggers
  import chex
  import jax
  import jax.numpy
  import numpy
  import optax
  import reverb
  import rlax
  import tree
label_desc: Use image node_id 6 for calling the MPOLearner obj's underlying member method code with example usage: obj._sgd_step(state, transitions) and returns: new_state,  metrics