docstr: Given a batch of inputs (e.g. states) and targets (e.g. discounted rewards), takes a training step: i.e. runs a
  forward pass and backward pass of the network, accumulates the gradients and applies an optimization step to
  update the weights.
  :param inputs: environment states (observation, etc.) as well extra inputs required by loss. Shape of ndarray
      is (batch_size, observation_space_size) or (batch_size, observation_space_size, stack_size)
  :param targets: targets required by  loss (e.g. sum of discounted rewards)
  :param scaler: value to scale gradients by before optimizing network weights
  :param additional_fetches: additional fetches to calculate and return. Each fetch is specified as (int, str)
      tuple of head-type-index and fetch-name. The tuple is obtained from each head.
  :param importance_weights: ndarray of shape (batch_size,) to multiply with batch loss.
  :return: tuple of total_loss, losses, norm_unclipped_grads, fetched_tensors
      total_loss (float): sum of all head losses
      losses (list of float): list of all losses. The order is list of target losses followed by list
          of regularization losses. The specifics of losses is dependant on the network parameters
          (number of heads, etc.)
      norm_unclippsed_grads (float): global norm of all gradients before any gradient clipping is applied
      fetched_tensors: all values for additional_fetches
docstr_tok: ["Given","a","batch","of","inputs","(","e.g",".","states",")","and","targets","(","e.g",".","discounted","rewards",")",",","takes","a","training","step",":","i.e",".","runs","a","forward","pass","and","backward","pass","of","the","network",",","accumulates","the","gradients","and","applies","an","optimization","step","to","update","the","weights",".",":","param","inputs",":","environment","states","(","observation",",","etc",".",")","as","well","extra","inputs","required","by","loss",".","Shape","of","ndarray","is","(","batch_size",",","observation_space_size",")","or","(","batch_size",",","observation_space_size",",","stack_size",")",":","param","targets",":","targets","required","by","loss","(","e.g",".","sum","of","discounted","rewards",")",":","param","scaler",":","value","to","scale","gradients","by","before","optimizing","network","weights",":","param","additional_fetches",":","additional","fetches","to","calculate","and","return",".","Each","fetch","is","specified","as","(","int",",","str",")","tuple","of","head-type-index","and","fetch-name",".","The","tuple","is","obtained","from","each","head",".",":","param","importance_weights",":","ndarray","of","shape","(","batch_size",",",")","to","multiply","with","batch","loss",".",":","return",":","tuple","of","total_loss",",","losses",",","norm_unclipped_grads",",","fetched_tensors","total_loss","(","float",")",":","sum","of","all","head","losses","losses","(","list","of","float",")",":","list","of","all","losses",".","The","order","is","list","of","target","losses","followed","by","list","of","regularization","losses",".","The","specifics","of","losses","is","dependant","on","the","network","parameters","(","number","of","heads",",","etc",".",")","norm_unclippsed_grads","(","float",")",":","global","norm","of","all","gradients","before","any","gradient","clipping","is","applied","fetched_tensors",":","all","values","for","additional_fetches"]
code: def train_on_batch(
      self,
      inputs: Dict[str, np.ndarray],
      targets: List[np.ndarray],
      scaler: float = 1.0,
      additional_fetches: list = None,
      importance_weights: np.ndarray = None,
  ) -> Tuple[float, List[float], float, list]:
      """
      Given a batch of inputs (e.g. states) and targets (e.g. discounted rewards), takes a training step: i.e. runs a
      forward pass and backward pass of the network, accumulates the gradients and applies an optimization step to
      update the weights.
      :param inputs: environment states (observation, etc.) as well extra inputs required by loss. Shape of ndarray
          is (batch_size, observation_space_size) or (batch_size, observation_space_size, stack_size)
      :param targets: targets required by  loss (e.g. sum of discounted rewards)
      :param scaler: value to scale gradients by before optimizing network weights
      :param additional_fetches: additional fetches to calculate and return. Each fetch is specified as (int, str)
          tuple of head-type-index and fetch-name. The tuple is obtained from each head.
      :param importance_weights: ndarray of shape (batch_size,) to multiply with batch loss.
      :return: tuple of total_loss, losses, norm_unclipped_grads, fetched_tensors
          total_loss (float): sum of all head losses
          losses (list of float): list of all losses. The order is list of target losses followed by list
              of regularization losses. The specifics of losses is dependant on the network parameters
              (number of heads, etc.)
          norm_unclippsed_grads (float): global norm of all gradients before any gradient clipping is applied
          fetched_tensors: all values for additional_fetches
      """
      loss = self.accumulate_gradients(
          inputs,
          targets,
          additional_fetches=additional_fetches,
          importance_weights=importance_weights,
      )
      self.apply_and_reset_gradients(self.accumulated_gradients, scaler)
      return loss
  
imports: import copy
  from typing import Any, Dict, Generator, List, Tuple, Union
  import numpy
  import mxnet
  from mxnet import autograd, gluon, nd
  from mxnet.ndarray import NDArray
  from rl_coach.architectures.architecture import Architecture
  from rl_coach.architectures.mxnet_components.heads.head import LOSS_OUT_TYPE_LOSS, LOSS_OUT_TYPE_REGULARIZATION
  from rl_coach.architectures.mxnet_components import utils
  from rl_coach.architectures.mxnet_components.savers import ParameterDictSaver, OnnxSaver
  from rl_coach.base_parameters import AgentParameters
  from rl_coach.logger import screen
  from rl_coach.saver import SaverCollection
  from rl_coach.spaces import SpacesDefinition
  from rl_coach.utils import force_list, squeeze_list
label_desc: Use image node_id 16 for calling the MxnetArchitecture obj's underlying member method code with example usage: obj.train_on_batch(inputs, targets, scaler, additional_fetches, importance_weights) and returns: loss