class_docstr: BC learner.
  
  This is the learning component of a BC agent. It takes a Transitions iterator
  as input and implements update functionality to learn from this iterator.
docstr: Behavior Cloning Learner.
  
  Args:
    networks: BC networks
    random_key: RNG key.
    loss_fn: BC loss to use.
    optimizer: Optax optimizer.
    prefetching_iterator: A sharded prefetching iterator as outputted from
      `acme.jax.utils.sharded_prefetch`. Please see the documentation for
      `sharded_prefetch` for more details.
    num_sgd_steps_per_step: Number of gradient updates per step.
    loss_has_aux: Whether the loss function returns auxiliary metrics as a
      second argument.
    logger: Logger.
    counter: Counter.
docstr_tok: ["Behavior","Cloning","Learner",".","Args",":","networks",":","BC","networks","random_key",":","RNG","key",".","loss_fn",":","BC","loss","to","use",".","optimizer",":","Optax","optimizer",".","prefetching_iterator",":","A","sharded","prefetching","iterator","as","outputted","from","`","acme.jax.utils.sharded_prefetch","`",".","Please","see","the","documentation","for","`","sharded_prefetch","`","for","more","details",".","num_sgd_steps_per_step",":","Number","of","gradient","updates","per","step",".","loss_has_aux",":","Whether","the","loss","function","returns","auxiliary","metrics","as","a","second","argument",".","logger",":","Logger",".","counter",":","Counter","."]
code: def __init__(
      self,
      networks: bc_networks.BCNetworks,
      random_key: networks_lib.PRNGKey,
      loss_fn: losses.BCLoss,
      optimizer: optax.GradientTransformation,
      prefetching_iterator: Iterator[types.Transition],
      num_sgd_steps_per_step: int,
      loss_has_aux: bool = False,
      logger: Optional[loggers.Logger] = None,
      counter: Optional[counting.Counter] = None,
  ):
      """Behavior Cloning Learner.
  
      Args:
        networks: BC networks
        random_key: RNG key.
        loss_fn: BC loss to use.
        optimizer: Optax optimizer.
        prefetching_iterator: A sharded prefetching iterator as outputted from
          `acme.jax.utils.sharded_prefetch`. Please see the documentation for
          `sharded_prefetch` for more details.
        num_sgd_steps_per_step: Number of gradient updates per step.
        loss_has_aux: Whether the loss function returns auxiliary metrics as a
          second argument.
        logger: Logger.
        counter: Counter.
      """
  
      def sgd_step(
          state: TrainingState,
          transitions: types.Transition,
      ) -> Tuple[TrainingState, Dict[str, jnp.ndarray]]:
          loss_and_grad = jax.value_and_grad(
              loss_fn, argnums=1, has_aux=loss_has_aux
          )
  
          # Compute losses and their gradients.
          key, key_input = jax.random.split(state.key)
          loss_result, gradients = loss_and_grad(
              networks, state.policy_params, key_input, transitions
          )
  
          # Combine the gradient across all devices (by taking their mean).
          gradients = jax.lax.pmean(
              gradients, axis_name=_PMAP_AXIS_NAME
          )
  
          # Compute and combine metrics across all devices.
          metrics = _create_loss_metrics(
              loss_has_aux, loss_result, gradients
          )
          metrics = jax.lax.pmean(metrics, axis_name=_PMAP_AXIS_NAME)
  
          policy_update, optimizer_state = optimizer.update(
              gradients, state.optimizer_state, state.policy_params
          )
          policy_params = optax.apply_updates(
              state.policy_params, policy_update
          )
  
          new_state = TrainingState(
              optimizer_state=optimizer_state,
              policy_params=policy_params,
              key=key,
              steps=state.steps + 1,
          )
  
          return new_state, metrics
  
      # General learner book-keeping and loggers.
      self._counter = counter or counting.Counter(prefix="learner")
      self._logger = logger or loggers.make_default_logger(
          "learner",
          asynchronous=True,
          serialize_fn=utils.fetch_devicearray,
          steps_key=self._counter.get_steps_key(),
      )
  
      # Split the input batch to `num_sgd_steps_per_step` minibatches in order
      # to achieve better performance on accelerators.
      sgd_step = utils.process_multiple_batches(
          sgd_step, num_sgd_steps_per_step
      )
      self._sgd_step = jax.pmap(sgd_step, axis_name=_PMAP_AXIS_NAME)
  
      random_key, init_key = jax.random.split(random_key)
      policy_params = networks.policy_network.init(init_key)
      optimizer_state = optimizer.init(policy_params)
  
      # Create initial state.
      state = TrainingState(
          optimizer_state=optimizer_state,
          policy_params=policy_params,
          key=random_key,
          steps=0,
      )
      self._state = utils.replicate_in_all_devices(state)
  
      self._timestamp = None
  
      self._prefetching_iterator = prefetching_iterator
  
imports: import time
  from typing import Dict, List, NamedTuple, Optional, Tuple, Union, Iterator
  import acme
  from acme import types
  from acme.agents.jax.bc import losses
  from acme.agents.jax.bc import networks
  from acme.jax import networks
  from acme.jax import utils
  from acme.utils import counting
  from acme.utils import loggers
  import jax
  import jax.numpy
  import optax
label_desc: Use image node_id 1 to create a new BCLearner object from inherited base classes: acme with example: obj = BCLearner(networks, random_key, loss_fn, optimizer, prefetching_iterator, num_sgd_steps_per_step, loss_has_aux, logger, counter)