code: def _orthogonal_regularization_loss(params: networks_lib.Params):
      """Orthogonal regularization.
  
      See equation (3) in https://arxiv.org/abs/1809.11096.
  
      Args:
        params: Dictionary of parameters to apply regualization for.
  
      Returns:
        A regularization loss term.
      """
      reg_loss = 0
      for key in params:
          if isinstance(params[key], Mapping):
              reg_loss += _orthogonal_regularization_loss(params[key])
              continue
          variable = params[key]
          assert len(variable.shape) in [1, 2, 4]
          if len(variable.shape) == 1:
              # This is a bias so do not apply regularization.
              continue
          if len(variable.shape) == 4:
              # CNN
              variable = jnp.reshape(variable, (-1, variable.shape[-1]))
          prod = jnp.matmul(jnp.transpose(variable), variable)
          reg_loss += jnp.sum(
              jnp.square(prod * (1 - jnp.eye(prod.shape[0])))
          )
      return reg_loss
  
imports: import functools
  import time
  from typing import Any, Dict, Iterator, List, Mapping, NamedTuple, Optional, Tuple
  import acme
  from acme import types
  from acme.agents.jax.value_dice 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
  import reverb
label_desc: Use image node_id 1 for calling a global function with example usage: _orthogonal_regularization_loss(params) and returns: reg_loss