class_docstr: AIRL Module.
class_docstr_tok: ["AIRL","Module","."]
code: def __call__(
      self,
      observations: networks_lib.Observation,
      actions: networks_lib.Action,
      next_observations: networks_lib.Observation,
      is_training: bool,
      rng: networks_lib.PRNGKey,
  ) -> networks_lib.Logits:
      g_output = DiscriminatorModule(
          environment_spec=self._environment_spec,
          use_action=self._use_action,
          use_next_obs=self._use_next_obs,
          network_core=self._g_core,
          observation_embedding=self._observation_embedding,
          name="airl_g",
      )(observations, actions, next_observations, is_training, rng)
      h_module = DiscriminatorModule(
          environment_spec=self._environment_spec,
          use_action=False,
          use_next_obs=False,
          network_core=self._h_core,
          observation_embedding=self._observation_embedding,
          name="airl_h",
      )
      return (
          g_output
          + self._discount
          * h_module(next_observations, (), (), is_training, rng)
          - h_module(observations, (), (), is_training, rng)
      )
  
imports: import dataclasses
  import functools
  from typing import Any, Callable, Generic, Iterable, Optional
  from acme import specs
  from acme import types
  from acme.jax import networks
  from acme.jax import utils
  from acme.jax.imitation_learning_types import DirectRLNetworks
  import haiku
  import jax
  from jax import numpy
  import numpy
label_desc: Use image node_id 2 for calling the AIRLModule obj's underlying member method code with example usage: obj.__call__(observations, actions, next_observations, is_training, rng) and returns: unknown