class_docstr: Regularized Q-learning.
  
  Implements DQNReg loss function: https://arxiv.org/abs/2101.03958.
  This is almost identical to QLearning except: 1) Adds a regularization term;
  2) Uses vanilla TD error without huber loss. 3) No reward clipping.
class_docstr_tok: ["Regularized","Q-learning",".","Implements","DQNReg","loss","function",":","https",":","\/\/arxiv.org\/abs\/2101.03958",".","This","is","almost","identical","to","QLearning","except",":","1",")","Adds","a","regularization","term",";","2",")","Uses","vanilla","TD","error","without","huber","loss",".","3",")","No","reward","clipping","."]
docstr: Calculate a loss on a single batch of data.
docstr_tok: ["Calculate","a","loss","on","a","single","batch","of","data","."]
code: def __call__(
      self,
      network: networks_lib.TypedFeedForwardNetwork,
      params: networks_lib.Params,
      target_params: networks_lib.Params,
      batch: reverb.ReplaySample,
      key: networks_lib.PRNGKey,
  ) -> Tuple[jax.Array, learning_lib.LossExtra]:
      """Calculate a loss on a single batch of data."""
      transitions: types.Transition = batch.data
  
      # Forward pass.
      key1, key2 = jax.random.split(key)
      q_tm1 = network.apply(
          params, transitions.observation, is_training=True, key=key1
      )
      q_t = network.apply(
          target_params,
          transitions.next_observation,
          is_training=True,
          key=key2,
      )
  
      d_t = (transitions.discount * self.discount).astype(jnp.float32)
  
      # Compute Q-learning TD-error.
      batch_error = jax.vmap(rlax.q_learning)
      td_error = batch_error(
          q_tm1, transitions.action, transitions.reward, d_t, q_t
      )
      td_error = 0.5 * jnp.square(td_error)
  
      def select(qtm1, action):
          return qtm1[action]
  
      q_regularizer = jax.vmap(select)(q_tm1, transitions.action)
  
      loss = self.regularizer_coeff * jnp.mean(
          q_regularizer
      ) + jnp.mean(td_error)
      extra = learning_lib.LossExtra(metrics={})
      return loss, extra
  
imports: import dataclasses
  from typing import Tuple
  from acme import types
  from acme.agents.jax.dqn import learning_lib
  from acme.jax import networks
  import chex
  import jax
  import jax.numpy
  import reverb
  import rlax
label_desc: Use image node_id 1 for calling the RegularizedQLearning obj's underlying member method code with example usage: obj.__call__(network, params, target_params, batch, key) and returns: loss,  extra, qtm1