class_docstr: AIRL Module.
class_docstr_tok: ["AIRL","Module","."]
code: def __init__(
      self,
      environment_spec: specs.EnvironmentSpec,
      use_action: bool,
      use_next_obs: bool,
      discount: float,
      g_core: Callable[..., Any],
      h_core: Callable[..., Any],
      observation_embedding: Callable[
          [networks_lib.Observation], jnp.ndarray
      ] = lambda x: x,
      name="airl",
  ):
      super().__init__(name=name)
      self._environment_spec = environment_spec
      self._use_action = use_action
      self._use_next_obs = use_next_obs
      self._discount = discount
      self._g_core = g_core
      self._h_core = h_core
      self._observation_embedding = observation_embedding
  
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 1 to create a new AIRLModule object from inherited base classes: hk with example: obj = AIRLModule(environment_spec, use_action, use_next_obs, discount, g_core, h_core, observation_embedding, name)